diff --git a/api/src/main/java/org/apache/flink/agents/api/EventType.java b/api/src/main/java/org/apache/flink/agents/api/EventType.java index 7724d45e81..971134ec33 100644 --- a/api/src/main/java/org/apache/flink/agents/api/EventType.java +++ b/api/src/main/java/org/apache/flink/agents/api/EventType.java @@ -42,16 +42,42 @@ public final class EventType { public static final String ContextRetrievalResponseEvent = org.apache.flink.agents.api.event.ContextRetrievalResponseEvent.EVENT_TYPE; + public static final String ShortTermWriteEvent = + org.apache.flink.agents.api.event.ShortTermWriteEvent.EVENT_TYPE; + public static final String ShortTermReadEvent = + org.apache.flink.agents.api.event.ShortTermReadEvent.EVENT_TYPE; + public static final String SensoryWriteEvent = + org.apache.flink.agents.api.event.SensoryWriteEvent.EVENT_TYPE; + public static final String SensoryReadEvent = + org.apache.flink.agents.api.event.SensoryReadEvent.EVENT_TYPE; + public static final String LongTermUpdateEvent = + org.apache.flink.agents.api.event.LongTermUpdateEvent.EVENT_TYPE; + public static final String LongTermGetEvent = + org.apache.flink.agents.api.event.LongTermGetEvent.EVENT_TYPE; + public static final String LongTermSearchEvent = + org.apache.flink.agents.api.event.LongTermSearchEvent.EVENT_TYPE; + + public static final String AgentRunBeginEvent = + org.apache.flink.agents.api.event.AgentRunBeginEvent.EVENT_TYPE; + private static final Map ALL_CONSTANTS = - Map.of( - "InputEvent", InputEvent, - "OutputEvent", OutputEvent, - "ChatRequestEvent", ChatRequestEvent, - "ChatResponseEvent", ChatResponseEvent, - "ToolRequestEvent", ToolRequestEvent, - "ToolResponseEvent", ToolResponseEvent, - "ContextRetrievalRequestEvent", ContextRetrievalRequestEvent, - "ContextRetrievalResponseEvent", ContextRetrievalResponseEvent); + Map.ofEntries( + Map.entry("InputEvent", InputEvent), + Map.entry("OutputEvent", OutputEvent), + Map.entry("ChatRequestEvent", ChatRequestEvent), + Map.entry("ChatResponseEvent", ChatResponseEvent), + Map.entry("ToolRequestEvent", ToolRequestEvent), + Map.entry("ToolResponseEvent", ToolResponseEvent), + Map.entry("ContextRetrievalRequestEvent", ContextRetrievalRequestEvent), + Map.entry("ContextRetrievalResponseEvent", ContextRetrievalResponseEvent), + Map.entry("ShortTermWriteEvent", ShortTermWriteEvent), + Map.entry("ShortTermReadEvent", ShortTermReadEvent), + Map.entry("SensoryWriteEvent", SensoryWriteEvent), + Map.entry("SensoryReadEvent", SensoryReadEvent), + Map.entry("LongTermUpdateEvent", LongTermUpdateEvent), + Map.entry("LongTermGetEvent", LongTermGetEvent), + Map.entry("LongTermSearchEvent", LongTermSearchEvent), + Map.entry("AgentRunBeginEvent", AgentRunBeginEvent)); /** * Returns the built-in event type constants as a name-to-value map for condition expressions. diff --git a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java index a0b9269e49..74e3cf9677 100644 --- a/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java +++ b/api/src/main/java/org/apache/flink/agents/api/agents/AgentExecutionOptions.java @@ -48,6 +48,10 @@ public class AgentExecutionOptions { public static final ConfigOption RAG_ASYNC = new ConfigOption<>("rag.async", Boolean.class, true); + /** Opt-in lifecycle event emitted at the beginning of each agent run. */ + public static final ConfigOption AGENT_RUN_BEGIN_EVENT = + new ConfigOption<>("agent-run.begin-event", Boolean.class, false); + /** Set to a positive value in milliseconds to enable short-term memory TTL; 0 disables it. */ public static final ConfigOption SHORT_TERM_MEMORY_STATE_TTL_MS = new ConfigOption<>("short-term-memory.state-ttl.ms", Long.class, 0L); diff --git a/api/src/main/java/org/apache/flink/agents/api/configuration/MemoryEventOptions.java b/api/src/main/java/org/apache/flink/agents/api/configuration/MemoryEventOptions.java new file mode 100644 index 0000000000..5cc31b013a --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/configuration/MemoryEventOptions.java @@ -0,0 +1,55 @@ +/* + * 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.configuration; + +/** + * Config options controlling memory observation events. + * + *

Resolution order per operation: the operation's own sub-key if explicitly configured, else the + * {@code memory.generate-event} master switch if explicitly configured, else the operation's + * built-in default (writes and all long-term ops: on; short-term/sensory reads: off). + */ +public class MemoryEventOptions { + + /** Master switch. Fallback for unset sub-switches; when unset itself, per-op defaults apply. */ + public static final ConfigOption MEMORY_GENERATE_EVENT = + new ConfigOption<>("memory.generate-event", Boolean.class, null); + + public static final ConfigOption SHORT_TERM_WRITE = + new ConfigOption<>("memory.generate-event.short-term-write", Boolean.class, null); + + public static final ConfigOption SHORT_TERM_READ = + new ConfigOption<>("memory.generate-event.short-term-read", Boolean.class, null); + + public static final ConfigOption SENSORY_WRITE = + new ConfigOption<>("memory.generate-event.sensory-write", Boolean.class, null); + + public static final ConfigOption SENSORY_READ = + new ConfigOption<>("memory.generate-event.sensory-read", Boolean.class, null); + + public static final ConfigOption LONG_TERM_UPDATE = + new ConfigOption<>("memory.generate-event.long-term-update", Boolean.class, null); + + public static final ConfigOption LONG_TERM_GET = + new ConfigOption<>("memory.generate-event.long-term-get", Boolean.class, null); + + public static final ConfigOption LONG_TERM_SEARCH = + new ConfigOption<>("memory.generate-event.long-term-search", Boolean.class, null); + + private MemoryEventOptions() {} +} diff --git a/api/src/main/java/org/apache/flink/agents/api/event/AgentRunBeginEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/AgentRunBeginEvent.java new file mode 100644 index 0000000000..90d74a389f --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/AgentRunBeginEvent.java @@ -0,0 +1,70 @@ +/* + * 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.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.flink.agents.api.Event; + +import java.util.Map; +import java.util.UUID; + +/** + * Lifecycle event marking the begin of one agent run. When enabled, it is emitted after an + * InputEvent arrives for a keyed partition and before any action of that run executes. Carries the + * textual form of the Flink key and the partition's short-term VALUE nodes as a dot-key flat map. + * Object nodes and empty-object structure are intentionally outside this event contract. + * + *

This event is opt-in through {@code agent-run.begin-event} (default false), independent of the + * {@code memory.generate-event} master switch. On the wire: {@code {id, type, attributes: {key, + * value}}}. + */ +public class AgentRunBeginEvent extends Event { + + public static final String EVENT_TYPE = "_agent_run_begin_event"; + + public AgentRunBeginEvent(String key, Map value) { + super(EVENT_TYPE, MemoryEvent.normalizeAttributes(key, value)); + } + + /** Deserialization constructor ({@code @JsonCreator}: see {@code MemoryEvent}'s note). */ + @JsonCreator + public AgentRunBeginEvent( + @JsonProperty("id") UUID id, + @JsonProperty("attributes") Map attributes) { + super(id, EVENT_TYPE, MemoryEvent.normalizeAttributes(attributes)); + } + + /** Converts a generic {@link Event} of this type into the typed view. */ + public static AgentRunBeginEvent fromEvent(Event event) { + return reconstructFrom(event, AgentRunBeginEvent::new); + } + + @JsonIgnore + public String getKey() { + return (String) getAttr("key"); + } + + @SuppressWarnings("unchecked") + @JsonIgnore + public Map getValue() { + return (Map) getAttr("value"); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/event/LongTermGetEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/LongTermGetEvent.java new file mode 100644 index 0000000000..7c6d09fa65 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/LongTermGetEvent.java @@ -0,0 +1,46 @@ +/* + * 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.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Map; +import java.util.UUID; + +/** + * Memory observation event: a long-term memory read by id. See {@link MemoryEvent} for the + * attribute contract. + */ +public class LongTermGetEvent extends MemoryEvent { + + public static final String EVENT_TYPE = "_long_term_get_event"; + + public LongTermGetEvent(String key, Map value) { + super(EVENT_TYPE, key, value); + } + + /** Deserialization constructor ({@code @JsonCreator}: see {@link MemoryEvent}'s note). */ + @JsonCreator + public LongTermGetEvent( + @JsonProperty("id") UUID id, + @JsonProperty("attributes") Map attributes) { + super(id, EVENT_TYPE, attributes); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/event/LongTermSearchEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/LongTermSearchEvent.java new file mode 100644 index 0000000000..3d5ddca69c --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/LongTermSearchEvent.java @@ -0,0 +1,56 @@ +/* + * 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.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.List; +import java.util.Map; +import java.util.UUID; + +/** + * Memory observation event: a semantic search over long-term memory. Unlike the other memory + * events, {@code value} maps each memory set to a map from query string to ordered hit list; {@link + * #getResults()} exposes that shape. + */ +public class LongTermSearchEvent extends MemoryEvent { + + public static final String EVENT_TYPE = "_long_term_search_event"; + + public LongTermSearchEvent(String key, Map value) { + super(EVENT_TYPE, key, value); + } + + /** Deserialization constructor ({@code @JsonCreator}: see {@link MemoryEvent}'s note). */ + @JsonCreator + public LongTermSearchEvent( + @JsonProperty("id") UUID id, + @JsonProperty("attributes") Map attributes) { + super(id, EVENT_TYPE, attributes); + } + + /** Typed view of {@link #getValue()}: memory set → query string → ordered hit list. */ + @SuppressWarnings("unchecked") + @JsonIgnore + public Map>>> getResults() { + return (Map>>>) (Map) getValue(); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/event/LongTermUpdateEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/LongTermUpdateEvent.java new file mode 100644 index 0000000000..00f1d556fb --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/LongTermUpdateEvent.java @@ -0,0 +1,88 @@ +/* + * 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.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +/** + * Memory observation event: a long-term memory update (adds and deletes). Item changes are grouped + * by memory set in {@code value}; {@code cleared_sets} identifies whole-set deletes. See {@link + * MemoryEvent} for the common attribute contract. + */ +public class LongTermUpdateEvent extends MemoryEvent { + + public static final String EVENT_TYPE = "_long_term_update_event"; + + public LongTermUpdateEvent(String key, Map value) { + this(key, value, Collections.emptyList()); + } + + public LongTermUpdateEvent(String key, Map value, List clearedSets) { + super(EVENT_TYPE, attributes(key, value, clearedSets)); + } + + /** Deserialization constructor ({@code @JsonCreator}: see {@link MemoryEvent}'s note). */ + @JsonCreator + public LongTermUpdateEvent( + @JsonProperty("id") UUID id, + @JsonProperty("attributes") Map attributes) { + super(id, EVENT_TYPE, normalizeUpdateAttributes(attributes)); + } + + private static Map attributes( + String key, Map value, List clearedSets) { + Map attributes = new LinkedHashMap<>(); + attributes.put("key", key); + attributes.put("value", value); + attributes.put( + "cleared_sets", + clearedSets == null ? Collections.emptyList() : List.copyOf(clearedSets)); + return attributes; + } + + private static Map normalizeUpdateAttributes(Map attributes) { + Map normalized = new LinkedHashMap<>(attributes); + Object clearedSets = normalized.get("cleared_sets"); + if (clearedSets == null) { + normalized.put("cleared_sets", Collections.emptyList()); + } else if (clearedSets instanceof List + && ((List) clearedSets).stream().allMatch(String.class::isInstance)) { + normalized.put("cleared_sets", List.copyOf((List) clearedSets)); + } else { + throw new IllegalArgumentException( + "Long-term update attribute 'cleared_sets' must be a list of strings."); + } + return normalized; + } + + @SuppressWarnings("unchecked") + @JsonIgnore + public List getClearedSets() { + Object clearedSets = getAttr("cleared_sets"); + return clearedSets instanceof List ? (List) clearedSets : Collections.emptyList(); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/event/MemoryEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/MemoryEvent.java new file mode 100644 index 0000000000..3618c4c348 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/MemoryEvent.java @@ -0,0 +1,134 @@ +/* + * 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}. + * + *

Attributes: {@code key} — the textual form of the Flink key the operation belongs to; {@code + * value} — the operation's folded JSON value map. On the wire the event serializes as {@code {id, + * type, attributes: {key, value}}}. + */ +public abstract class MemoryEvent extends Event { + + private static final ObjectMapper MAPPER = + JsonMapper.builder().disable(JsonWriteFeature.WRITE_NAN_AS_STRINGS).build(); + + /** + * The seven memory event types, each mapped to its deserialization constructor ({@code (id, + * attributes)}). Single source of truth for both {@link #isMemoryType} and {@link #fromEvent}. + */ + private static final Map, MemoryEvent>> FACTORIES = + Map., MemoryEvent>>of( + ShortTermWriteEvent.EVENT_TYPE, ShortTermWriteEvent::new, + ShortTermReadEvent.EVENT_TYPE, ShortTermReadEvent::new, + SensoryWriteEvent.EVENT_TYPE, SensoryWriteEvent::new, + SensoryReadEvent.EVENT_TYPE, SensoryReadEvent::new, + LongTermUpdateEvent.EVENT_TYPE, LongTermUpdateEvent::new, + LongTermGetEvent.EVENT_TYPE, LongTermGetEvent::new, + LongTermSearchEvent.EVENT_TYPE, LongTermSearchEvent::new); + + /** Returns true iff {@code type} is one of the seven memory operation event types. */ + public static boolean isMemoryType(String type) { + return type != null && FACTORIES.containsKey(type); + } + + protected MemoryEvent(String type, String key, Map value) { + super(type, normalizeAttributes(key, value)); + } + + protected MemoryEvent(String type, Map attributes) { + super(type, normalizeAttributes(attributes)); + } + + /** + * Deserialization constructor path. The subclass {@code @JsonCreator}s are REQUIRED: instances + * land in ActionState.outputEvents and deserialize polymorphically via the {@code @class} mixin + * in ActionStateSerde — without an explicit creator Jackson cannot instantiate the subtype. + */ + protected MemoryEvent(UUID id, String type, Map attributes) { + super(id, type, normalizeAttributes(attributes)); + } + + static Map normalizeAttributes(Map attributes) { + if (attributes == null) { + throw new IllegalArgumentException("Memory observation attributes must not be null."); + } + Object key = attributes.get("key"); + if (!(key instanceof String)) { + throw new IllegalArgumentException( + "Memory observation attribute 'key' must be a string."); + } + Object value = attributes.get("value"); + if (!(value instanceof Map)) { + throw new IllegalArgumentException( + "Memory observation attribute 'value' must be a map."); + } + Map normalized = new LinkedHashMap<>(attributes); + try { + normalized.put( + "value", MAPPER.readValue(MAPPER.writeValueAsBytes(value), Object.class)); + } catch (IOException | RuntimeException e) { + throw new IllegalArgumentException( + "Memory observation value must be JSON serializable.", e); + } + return normalized; + } + + static Map normalizeAttributes(String key, Map value) { + Map attributes = new LinkedHashMap<>(); + attributes.put("key", key); + attributes.put("value", value); + return normalizeAttributes(attributes); + } + + /** Converts a generic {@link Event} carrying a memory type into its typed subclass view. */ + public static MemoryEvent fromEvent(Event event) { + BiFunction, MemoryEvent> factory = FACTORIES.get(event.getType()); + if (factory == null) { + throw new IllegalArgumentException("Not a memory event type: " + event.getType()); + } + return reconstructFrom(event, factory); + } + + @JsonIgnore + public String getKey() { + return (String) getAttr("key"); + } + + @SuppressWarnings("unchecked") + @JsonIgnore + public Map getValue() { + return (Map) getAttr("value"); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/event/SensoryReadEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/SensoryReadEvent.java new file mode 100644 index 0000000000..a077aa07d8 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/SensoryReadEvent.java @@ -0,0 +1,46 @@ +/* + * 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.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Map; +import java.util.UUID; + +/** + * Memory observation event: a read from sensory memory. See {@link MemoryEvent} for the attribute + * contract. + */ +public class SensoryReadEvent extends MemoryEvent { + + public static final String EVENT_TYPE = "_sensory_read_event"; + + public SensoryReadEvent(String key, Map value) { + super(EVENT_TYPE, key, value); + } + + /** Deserialization constructor ({@code @JsonCreator}: see {@link MemoryEvent}'s note). */ + @JsonCreator + public SensoryReadEvent( + @JsonProperty("id") UUID id, + @JsonProperty("attributes") Map attributes) { + super(id, EVENT_TYPE, attributes); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/event/SensoryWriteEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/SensoryWriteEvent.java new file mode 100644 index 0000000000..6b420d86d9 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/SensoryWriteEvent.java @@ -0,0 +1,46 @@ +/* + * 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.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Map; +import java.util.UUID; + +/** + * Memory observation event: a write to sensory memory. See {@link MemoryEvent} for the attribute + * contract. + */ +public class SensoryWriteEvent extends MemoryEvent { + + public static final String EVENT_TYPE = "_sensory_write_event"; + + public SensoryWriteEvent(String key, Map value) { + super(EVENT_TYPE, key, value); + } + + /** Deserialization constructor ({@code @JsonCreator}: see {@link MemoryEvent}'s note). */ + @JsonCreator + public SensoryWriteEvent( + @JsonProperty("id") UUID id, + @JsonProperty("attributes") Map attributes) { + super(id, EVENT_TYPE, attributes); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ShortTermReadEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ShortTermReadEvent.java new file mode 100644 index 0000000000..8c2be4cfa0 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/ShortTermReadEvent.java @@ -0,0 +1,46 @@ +/* + * 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.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Map; +import java.util.UUID; + +/** + * Memory observation event: a read from short-term memory. See {@link MemoryEvent} for the + * attribute contract. + */ +public class ShortTermReadEvent extends MemoryEvent { + + public static final String EVENT_TYPE = "_short_term_read_event"; + + public ShortTermReadEvent(String key, Map value) { + super(EVENT_TYPE, key, value); + } + + /** Deserialization constructor ({@code @JsonCreator}: see {@link MemoryEvent}'s note). */ + @JsonCreator + public ShortTermReadEvent( + @JsonProperty("id") UUID id, + @JsonProperty("attributes") Map attributes) { + super(id, EVENT_TYPE, attributes); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/event/ShortTermWriteEvent.java b/api/src/main/java/org/apache/flink/agents/api/event/ShortTermWriteEvent.java new file mode 100644 index 0000000000..bc7b79a8fb --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/ShortTermWriteEvent.java @@ -0,0 +1,46 @@ +/* + * 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.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Map; +import java.util.UUID; + +/** + * Memory observation event: a write to short-term memory. See {@link MemoryEvent} for the attribute + * contract. + */ +public class ShortTermWriteEvent extends MemoryEvent { + + public static final String EVENT_TYPE = "_short_term_write_event"; + + public ShortTermWriteEvent(String key, Map value) { + super(EVENT_TYPE, key, value); + } + + /** Deserialization constructor ({@code @JsonCreator}: see {@link MemoryEvent}'s note). */ + @JsonCreator + public ShortTermWriteEvent( + @JsonProperty("id") UUID id, + @JsonProperty("attributes") Map attributes) { + super(id, EVENT_TYPE, attributes); + } +} diff --git a/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java b/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java index 5ae8bc7e8c..b09e6a13a9 100644 --- a/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/CrossLanguageEventSnapshotTest.java @@ -23,10 +23,13 @@ import org.apache.flink.agents.api.agents.OutputSchema; import org.apache.flink.agents.api.chat.messages.ChatMessage; import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.event.AgentRunBeginEvent; import org.apache.flink.agents.api.event.ChatRequestEvent; import org.apache.flink.agents.api.event.ChatResponseEvent; import org.apache.flink.agents.api.event.ContextRetrievalRequestEvent; import org.apache.flink.agents.api.event.ContextRetrievalResponseEvent; +import org.apache.flink.agents.api.event.MemoryEvent; +import org.apache.flink.agents.api.event.ShortTermWriteEvent; import org.apache.flink.agents.api.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; import org.apache.flink.agents.api.tools.ToolResponse; @@ -447,6 +450,75 @@ void javaCanDeserializeContextRetrievalResponseEventFromPythonSnapshot() throws assertEquals("doc-1", docs.get(0).getId()); } + // ── Memory observation events ────────────────────────────────────────── + + private static ShortTermWriteEvent buildShortTermWriteEvent() { + Map value = new LinkedHashMap<>(); + value.put("user.tier", "gold"); + value.put("profile.m_a01", null); + Map attrs = new LinkedHashMap<>(); + attrs.put("key", "user-42"); + attrs.put("value", value); + return new ShortTermWriteEvent(FIXED_EVENT_ID, attrs); + } + + @Test + void regenerateShortTermWriteEventJavaSnapshot() throws Exception { + assumeTrue(regenerateRequested(), "Set -Dregenerate.snapshots=true to refresh."); + writeJavaSnapshot("short_term_write_event.json", buildShortTermWriteEvent()); + } + + @Test + void shortTermWriteEventJavaSnapshotIsStable() throws Exception { + assertJavaSnapshotStable("short_term_write_event.json", buildShortTermWriteEvent()); + } + + @Test + void javaCanDeserializeShortTermWriteEventFromPythonSnapshot() throws Exception { + Event base = readPythonSnapshot("short_term_write_event.json"); + ShortTermWriteEvent typed = (ShortTermWriteEvent) MemoryEvent.fromEvent(base); + + assertEquals(FIXED_EVENT_ID, typed.getId()); + assertEquals("user-42", typed.getKey()); + assertEquals("gold", typed.getValue().get("user.tier")); + assertTrue(typed.getValue().containsKey("profile.m_a01")); + assertEquals(null, typed.getValue().get("profile.m_a01")); + } + + // ── Agent-run lifecycle events ───────────────────────────────────────── + + private static AgentRunBeginEvent buildAgentRunBeginEvent() { + Map value = new LinkedHashMap<>(); + value.put("user.tier", "gold"); + value.put("user.address.city", "SF"); + Map attrs = new LinkedHashMap<>(); + attrs.put("key", "user-42"); + attrs.put("value", value); + return new AgentRunBeginEvent(FIXED_EVENT_ID, attrs); + } + + @Test + void regenerateAgentRunBeginEventJavaSnapshot() throws Exception { + assumeTrue(regenerateRequested(), "Set -Dregenerate.snapshots=true to refresh."); + writeJavaSnapshot("agent_run_begin_event.json", buildAgentRunBeginEvent()); + } + + @Test + void agentRunBeginEventJavaSnapshotIsStable() throws Exception { + assertJavaSnapshotStable("agent_run_begin_event.json", buildAgentRunBeginEvent()); + } + + @Test + void javaCanDeserializeAgentRunBeginEventFromPythonSnapshot() throws Exception { + Event base = readPythonSnapshot("agent_run_begin_event.json"); + AgentRunBeginEvent typed = AgentRunBeginEvent.fromEvent(base); + + assertEquals(FIXED_EVENT_ID, typed.getId()); + assertEquals("user-42", typed.getKey()); + assertEquals("gold", typed.getValue().get("user.tier")); + assertEquals("SF", typed.getValue().get("user.address.city")); + } + // ── Generic Event with primitive attributes (user-authored axis) ─────── private static final String GENERIC_EVENT_TYPE = "_my_custom_event"; diff --git a/api/src/test/java/org/apache/flink/agents/api/EventTypeTest.java b/api/src/test/java/org/apache/flink/agents/api/EventTypeTest.java index 4b7a07b18f..8aca070de5 100644 --- a/api/src/test/java/org/apache/flink/agents/api/EventTypeTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/EventTypeTest.java @@ -18,10 +18,18 @@ package org.apache.flink.agents.api; +import org.apache.flink.agents.api.event.AgentRunBeginEvent; import org.apache.flink.agents.api.event.ChatRequestEvent; import org.apache.flink.agents.api.event.ChatResponseEvent; import org.apache.flink.agents.api.event.ContextRetrievalRequestEvent; import org.apache.flink.agents.api.event.ContextRetrievalResponseEvent; +import org.apache.flink.agents.api.event.LongTermGetEvent; +import org.apache.flink.agents.api.event.LongTermSearchEvent; +import org.apache.flink.agents.api.event.LongTermUpdateEvent; +import org.apache.flink.agents.api.event.SensoryReadEvent; +import org.apache.flink.agents.api.event.SensoryWriteEvent; +import org.apache.flink.agents.api.event.ShortTermReadEvent; +import org.apache.flink.agents.api.event.ShortTermWriteEvent; import org.apache.flink.agents.api.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; import org.junit.jupiter.api.Test; @@ -37,15 +45,27 @@ class EventTypeTest { @Test void allConstantsProvidesAnUnmodifiableNameToValueMap() { assertEquals( - Map.of( - "InputEvent", InputEvent.EVENT_TYPE, - "OutputEvent", OutputEvent.EVENT_TYPE, - "ChatRequestEvent", ChatRequestEvent.EVENT_TYPE, - "ChatResponseEvent", ChatResponseEvent.EVENT_TYPE, - "ToolRequestEvent", ToolRequestEvent.EVENT_TYPE, - "ToolResponseEvent", ToolResponseEvent.EVENT_TYPE, - "ContextRetrievalRequestEvent", ContextRetrievalRequestEvent.EVENT_TYPE, - "ContextRetrievalResponseEvent", ContextRetrievalResponseEvent.EVENT_TYPE), + Map.ofEntries( + Map.entry("InputEvent", InputEvent.EVENT_TYPE), + Map.entry("OutputEvent", OutputEvent.EVENT_TYPE), + Map.entry("ChatRequestEvent", ChatRequestEvent.EVENT_TYPE), + Map.entry("ChatResponseEvent", ChatResponseEvent.EVENT_TYPE), + Map.entry("ToolRequestEvent", ToolRequestEvent.EVENT_TYPE), + Map.entry("ToolResponseEvent", ToolResponseEvent.EVENT_TYPE), + Map.entry( + "ContextRetrievalRequestEvent", + ContextRetrievalRequestEvent.EVENT_TYPE), + Map.entry( + "ContextRetrievalResponseEvent", + ContextRetrievalResponseEvent.EVENT_TYPE), + Map.entry("ShortTermWriteEvent", ShortTermWriteEvent.EVENT_TYPE), + Map.entry("ShortTermReadEvent", ShortTermReadEvent.EVENT_TYPE), + Map.entry("SensoryWriteEvent", SensoryWriteEvent.EVENT_TYPE), + Map.entry("SensoryReadEvent", SensoryReadEvent.EVENT_TYPE), + Map.entry("LongTermUpdateEvent", LongTermUpdateEvent.EVENT_TYPE), + Map.entry("LongTermGetEvent", LongTermGetEvent.EVENT_TYPE), + Map.entry("LongTermSearchEvent", LongTermSearchEvent.EVENT_TYPE), + Map.entry("AgentRunBeginEvent", AgentRunBeginEvent.EVENT_TYPE)), EventType.allConstants()); assertThrows( UnsupportedOperationException.class, diff --git a/api/src/test/java/org/apache/flink/agents/api/event/AgentRunBeginEventTest.java b/api/src/test/java/org/apache/flink/agents/api/event/AgentRunBeginEventTest.java new file mode 100644 index 0000000000..57cb358eba --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/event/AgentRunBeginEventTest.java @@ -0,0 +1,151 @@ +/* + * 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.databind.ObjectMapper; +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.EventType; +import org.junit.jupiter.api.Test; + +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** Tests for {@link AgentRunBeginEvent}. */ +public class AgentRunBeginEventTest { + + @Test + void testTypeAndAggregateConstant() { + assertEquals("_agent_run_begin_event", AgentRunBeginEvent.EVENT_TYPE); + assertEquals(AgentRunBeginEvent.EVENT_TYPE, EventType.AgentRunBeginEvent); + // lifecycle event, NOT a memory operation type (no observation suppression) + assertFalse(MemoryEvent.isMemoryType(AgentRunBeginEvent.EVENT_TYPE)); + } + + @Test + void testCarriesKeyAndStmValues() { + Map stm = new LinkedHashMap<>(); + stm.put("user.tier", "gold"); + stm.put("user.address.city", "SF"); + AgentRunBeginEvent event = new AgentRunBeginEvent("user-42", stm); + + assertEquals("user-42", event.getKey()); + assertEquals(stm, event.getValue()); + assertEquals("user-42", event.getAttributes().get("key")); + } + + @Test + void testLiveConstructorNormalizesValuesToJsonEquivalentObjects() { + Map stm = new LinkedHashMap<>(); + stm.put("bytes", new byte[] {1, 2, 3}); + stm.put("pojo", new MemoryEventTest.ObservationValuePojo("hello", 7)); + + AgentRunBeginEvent event = new AgentRunBeginEvent("user-42", stm); + + assertEquals("AQID", event.getValue().get("bytes")); + assertEquals(Map.of("name", "hello", "count", 7), event.getValue().get("pojo")); + assertInstanceOf(byte[].class, stm.get("bytes")); + assertInstanceOf(MemoryEventTest.ObservationValuePojo.class, stm.get("pojo")); + } + + @Test + void testLiveConstructorRejectsPojoWithNonFiniteNumber() { + Map stm = + Map.of("pojo", new MemoryEventTest.NonFiniteValuePojo(Double.POSITIVE_INFINITY)); + + assertThrows(IllegalArgumentException.class, () -> new AgentRunBeginEvent("user-42", stm)); + } + + @Test + void testFromEventNormalizesRawValue() { + Map stm = new LinkedHashMap<>(); + stm.put("pojo", new MemoryEventTest.ObservationValuePojo("hello", 7)); + Map attributes = new LinkedHashMap<>(); + attributes.put("key", "k"); + attributes.put("value", stm); + Event generic = new Event(AgentRunBeginEvent.EVENT_TYPE, attributes); + + AgentRunBeginEvent restored = AgentRunBeginEvent.fromEvent(generic); + + assertEquals("k", restored.getKey()); + assertEquals(Map.of("name", "hello", "count", 7), restored.getValue().get("pojo")); + assertInstanceOf(MemoryEventTest.ObservationValuePojo.class, stm.get("pojo")); + } + + @Test + void testFromEventPreservesFrameworkMetadata() { + UUID upstreamEventId = UUID.randomUUID(); + Event generic = + new Event( + AgentRunBeginEvent.EVENT_TYPE, + Map.of("key", "k", "value", Map.of("profile.name", "Alice"))); + generic.setSourceTimestamp(123456789L); + generic.setUpstreamEventId(upstreamEventId); + generic.setUpstreamActionName("start_agent_run"); + + AgentRunBeginEvent restored = AgentRunBeginEvent.fromEvent(generic); + + assertEquals(generic.getId(), restored.getId()); + assertEquals(123456789L, restored.getSourceTimestamp()); + assertEquals(upstreamEventId, restored.getUpstreamEventId()); + assertEquals("start_agent_run", restored.getUpstreamActionName()); + } + + @Test + void testFromEventRejectsRawNonFiniteNumber() { + Map attributes = new LinkedHashMap<>(); + attributes.put("key", "k"); + attributes.put("value", Map.of("nested", List.of(Float.NEGATIVE_INFINITY))); + Event generic = new Event(AgentRunBeginEvent.EVENT_TYPE, attributes); + + assertThrows(IllegalArgumentException.class, () -> AgentRunBeginEvent.fromEvent(generic)); + } + + @Test + void testFromEventRejectsIncompleteAttributes() { + Event missingValue = new Event(AgentRunBeginEvent.EVENT_TYPE, Map.of("key", "user-42")); + + assertThrows( + IllegalArgumentException.class, () -> AgentRunBeginEvent.fromEvent(missingValue)); + } + + @Test + void testJsonSerdeRoundTrip() throws Exception { + // Exercises the typed @JsonCreator(id, attributes) path via direct JSON roundtrip. + // ActionStateSerdeTest covers the real durable-state restoration path. + Map stm = new LinkedHashMap<>(); + stm.put("user.tier", "gold"); + stm.put("user.address.city", "SF"); + AgentRunBeginEvent original = new AgentRunBeginEvent("user-42", stm); + + ObjectMapper mapper = new ObjectMapper(); + String json = mapper.writeValueAsString(original); + AgentRunBeginEvent restored = mapper.readValue(json, AgentRunBeginEvent.class); + + assertEquals(original.getId(), restored.getId()); + assertEquals("user-42", restored.getKey()); + assertEquals(stm, restored.getValue()); + } +} diff --git a/api/src/test/java/org/apache/flink/agents/api/event/MemoryEventTest.java b/api/src/test/java/org/apache/flink/agents/api/event/MemoryEventTest.java new file mode 100644 index 0000000000..7747ef68fb --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/event/MemoryEventTest.java @@ -0,0 +1,279 @@ +/* + * 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 org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.EventType; +import org.apache.flink.agents.api.InputEvent; +import org.junit.jupiter.api.Test; + +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Tests for {@link MemoryEvent} and its seven concrete subclasses. */ +public class MemoryEventTest { + + @Test + void testSubclassesPinTypes() { + assertEquals("_short_term_write_event", ShortTermWriteEvent.EVENT_TYPE); + assertEquals("_short_term_read_event", ShortTermReadEvent.EVENT_TYPE); + assertEquals("_sensory_write_event", SensoryWriteEvent.EVENT_TYPE); + assertEquals("_sensory_read_event", SensoryReadEvent.EVENT_TYPE); + assertEquals("_long_term_update_event", LongTermUpdateEvent.EVENT_TYPE); + assertEquals("_long_term_get_event", LongTermGetEvent.EVENT_TYPE); + assertEquals("_long_term_search_event", LongTermSearchEvent.EVENT_TYPE); + + assertEquals( + ShortTermWriteEvent.EVENT_TYPE, + new ShortTermWriteEvent("k", new LinkedHashMap<>()).getType()); + assertEquals( + LongTermSearchEvent.EVENT_TYPE, + new LongTermSearchEvent("k", new LinkedHashMap<>()).getType()); + } + + @Test + void testIsMemoryType() { + assertTrue(MemoryEvent.isMemoryType("_short_term_write_event")); + assertTrue(MemoryEvent.isMemoryType("_short_term_read_event")); + assertTrue(MemoryEvent.isMemoryType("_sensory_write_event")); + assertTrue(MemoryEvent.isMemoryType("_sensory_read_event")); + assertTrue(MemoryEvent.isMemoryType("_long_term_update_event")); + assertTrue(MemoryEvent.isMemoryType("_long_term_get_event")); + assertTrue(MemoryEvent.isMemoryType("_long_term_search_event")); + assertFalse(MemoryEvent.isMemoryType("_input_event")); + assertFalse(MemoryEvent.isMemoryType("_agent_run_begin_event")); + assertFalse(MemoryEvent.isMemoryType(null)); + } + + @Test + void testEventTypeAggregateConstants() { + assertEquals(ShortTermWriteEvent.EVENT_TYPE, EventType.ShortTermWriteEvent); + assertEquals(ShortTermReadEvent.EVENT_TYPE, EventType.ShortTermReadEvent); + assertEquals(SensoryWriteEvent.EVENT_TYPE, EventType.SensoryWriteEvent); + assertEquals(SensoryReadEvent.EVENT_TYPE, EventType.SensoryReadEvent); + assertEquals(LongTermUpdateEvent.EVENT_TYPE, EventType.LongTermUpdateEvent); + assertEquals(LongTermGetEvent.EVENT_TYPE, EventType.LongTermGetEvent); + assertEquals(LongTermSearchEvent.EVENT_TYPE, EventType.LongTermSearchEvent); + } + + @Test + void testFromEventRejectsNonMemoryType() { + Event nonMemory = new InputEvent("data"); + assertThrows(IllegalArgumentException.class, () -> MemoryEvent.fromEvent(nonMemory)); + } + + @Test + void testKeyValueLiveInAttributes() { + Map value = new LinkedHashMap<>(); + value.put("user.tier", "gold"); + MemoryEvent event = new ShortTermWriteEvent("user-42", value); + + assertEquals("user-42", event.getKey()); + assertEquals("gold", event.getValue().get("user.tier")); + assertEquals("user-42", event.getAttributes().get("key")); + assertEquals(value, event.getAttributes().get("value")); + } + + @Test + void testLiveConstructorNormalizesValuesToJsonEquivalentObjects() { + Map value = new LinkedHashMap<>(); + value.put("bytes", new byte[] {1, 2, 3}); + value.put("pojo", new ObservationValuePojo("hello", 7)); + + MemoryEvent event = new ShortTermWriteEvent("user-42", value); + + assertEquals("AQID", event.getValue().get("bytes")); + assertEquals(Map.of("name", "hello", "count", 7), event.getValue().get("pojo")); + assertInstanceOf(byte[].class, value.get("bytes")); + assertInstanceOf(ObservationValuePojo.class, value.get("pojo")); + } + + @Test + void testLiveConstructorRejectsNonJsonSerializableValue() { + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> new ShortTermWriteEvent("user-42", Map.of("bad", new Object()))); + + assertTrue( + error.getMessage().contains("Memory observation value must be JSON serializable")); + } + + @Test + void testLiveConstructorRejectsNestedNonFiniteNumbers() { + for (Number nonFinite : + List.of( + Double.NaN, + Double.POSITIVE_INFINITY, + Double.NEGATIVE_INFINITY, + Float.NaN, + Float.POSITIVE_INFINITY, + Float.NEGATIVE_INFINITY)) { + Map value = Map.of("nested", Map.of("numbers", List.of(nonFinite))); + + assertThrows( + IllegalArgumentException.class, + () -> new ShortTermWriteEvent("user-42", value), + nonFinite.toString()); + } + } + + @Test + void testFromEventDispatchesToSubclassAndNormalizesRawValue() { + Map value = new LinkedHashMap<>(); + value.put("bytes", new byte[] {1, 2, 3}); + Map attributes = new LinkedHashMap<>(); + attributes.put("key", "k1"); + attributes.put("value", value); + Event generic = new Event(LongTermUpdateEvent.EVENT_TYPE, attributes); + + MemoryEvent restored = MemoryEvent.fromEvent(generic); + + assertInstanceOf(LongTermUpdateEvent.class, restored); + assertEquals(LongTermUpdateEvent.EVENT_TYPE, restored.getType()); + assertEquals("k1", restored.getKey()); + assertEquals("AQID", restored.getValue().get("bytes")); + assertInstanceOf(byte[].class, value.get("bytes")); + } + + @Test + void testFromEventPreservesFrameworkMetadata() { + UUID upstreamEventId = UUID.randomUUID(); + Event generic = + new Event( + LongTermUpdateEvent.EVENT_TYPE, + Map.of("key", "k1", "value", Map.of("profile.name", "Alice"))); + generic.setSourceTimestamp(123456789L); + generic.setUpstreamEventId(upstreamEventId); + generic.setUpstreamActionName("update_profile"); + + MemoryEvent restored = MemoryEvent.fromEvent(generic); + + assertEquals(generic.getId(), restored.getId()); + assertEquals(123456789L, restored.getSourceTimestamp()); + assertEquals(upstreamEventId, restored.getUpstreamEventId()); + assertEquals("update_profile", restored.getUpstreamActionName()); + } + + @Test + void testFromEventRejectsRawNonFiniteNumber() { + Map attributes = new LinkedHashMap<>(); + attributes.put("key", "k1"); + attributes.put("value", Map.of("nested", List.of(Double.NaN))); + Event generic = new Event(LongTermUpdateEvent.EVENT_TYPE, attributes); + + assertThrows(IllegalArgumentException.class, () -> MemoryEvent.fromEvent(generic)); + } + + @Test + void testConstructorsRejectIncompleteAttributes() { + assertThrows( + IllegalArgumentException.class, + () -> new ShortTermWriteEvent((String) null, new LinkedHashMap<>())); + assertThrows(IllegalArgumentException.class, () -> new ShortTermWriteEvent("k", null)); + + Event missingValue = new Event(ShortTermWriteEvent.EVENT_TYPE, Map.of("key", "user-42")); + assertThrows(IllegalArgumentException.class, () -> MemoryEvent.fromEvent(missingValue)); + } + + @Test + void testLongTermSearchTypedResults() { + Map value = new LinkedHashMap<>(); + value.put( + "policies", + Map.of( + "refund policy", + List.of( + Map.of("id", "p_01", "value", "7-day refund", "score", 0.92), + Map.of( + "id", + "p_02", + "value", + "no custom refunds", + "score", + 0.81)))); + LongTermSearchEvent event = new LongTermSearchEvent("user-42", value); + + Map>>> results = event.getResults(); + assertEquals(2, results.get("policies").get("refund policy").size()); + assertEquals("p_01", results.get("policies").get("refund policy").get(0).get("id")); + } + + @Test + void testLongTermUpdateCarriesClearedSets() { + LongTermUpdateEvent event = + new LongTermUpdateEvent( + "k", Map.of("prefs", Map.of("m2", "new")), List.of("prefs")); + + assertEquals(List.of("prefs"), event.getClearedSets()); + assertEquals(Map.of("m2", "new"), event.getValue().get("prefs")); + } + + public static class ObservationValuePojo { + private final String name; + private final int count; + + public ObservationValuePojo(String name, int count) { + this.name = name; + this.count = count; + } + + public String getName() { + return name; + } + + public int getCount() { + return count; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof ObservationValuePojo)) return false; + ObservationValuePojo that = (ObservationValuePojo) o; + return count == that.count && Objects.equals(name, that.name); + } + + @Override + public int hashCode() { + return Objects.hash(name, count); + } + } + + public static class NonFiniteValuePojo { + private final double value; + + public NonFiniteValuePojo(double value) { + this.value = value; + } + + public double getValue() { + return value; + } + } +} diff --git a/docs/content/docs/development/memory/long_term_memory.md b/docs/content/docs/development/memory/long_term_memory.md index ef0adc8ba1..666c697b9e 100644 --- a/docs/content/docs/development/memory/long_term_memory.md +++ b/docs/content/docs/development/memory/long_term_memory.md @@ -539,7 +539,7 @@ 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 automatically provides context isolation through Flink's keyed partition model. Each keyed partition normally maintains its own isolated set of memories. The isolation hierarchy works as follows: - **Job-level** (`JOB_IDENTIFIER`): Separates memories between different Flink jobs @@ -548,4 +548,4 @@ The isolation hierarchy works as follows: This means you can reuse the same memory set name across different partitions, and each partition will normally access only its own memories. -> **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. +> **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. diff --git a/docs/content/docs/development/memory/memory_events.md b/docs/content/docs/development/memory/memory_events.md new file mode 100644 index 0000000000..c89235012f --- /dev/null +++ b/docs/content/docs/development/memory/memory_events.md @@ -0,0 +1,246 @@ +--- +title: Memory Events +weight: 5 +type: docs +--- + + +## Overview + +Flink Agents can publish memory operations as framework-generated events. These events make it possible to audit what an action read or wrote, subscribe another action to memory changes, and inspect short-lived memory values from the Event Log for debugging. + +When an action finishes, the framework folds the memory operations performed by that action into **memory events**. It emits at most one event for each memory scope and operation kind, for example one short-term write event and one sensory write event. Failed or unfinished actions do not emit memory events. An action failure fails the operator task, and recovery rebuilds the interpreter instead of continuing with later actions in the same interpreter instance. + +Memory events use the same event infrastructure as other Flink Agents events. They can be written to the [Event Log]({{< ref "docs/operations/monitoring#event-log" >}}), observed by registered `EventListener`s, and used to trigger actions through normal event routing. + +The framework can also emit an **agent-run begin event** (`_agent_run_begin_event`) when an `InputEvent` starts a run for a key. This opt-in event is disabled by default. When enabled, it is emitted before any action in that run executes and carries the key's short-term memory **values** as a flat map. Its Event Lineage points to the triggering `InputEvent` through the framework virtual Action `agent_run_begin_action`, so the event and any Actions it triggers remain in the same Trace Tree. Object nodes and empty-object structure are outside the event contract. + +## When to Use + +Memory events are useful when you need to: + +- audit which memory entries an action changed; +- trigger follow-up actions from memory changes; +- inspect memory reads during debugging by enabling read events; +- inspect short-term or sensory values from an Event Log captured at `VERBOSE` level. + +Memory events are not a replacement for the memory backends themselves. In particular, long-term memory events describe requested operations and observed read/search results; they do not provide a complete copy of the long-term memory store. + +## Emission Model + +Memory events are emitted at the action finish boundary, together with the user events produced by that action. The framework records memory activity while the action runs, folds the records by event type, and then appends the generated events to the action's output events. + +Framework observation is best-effort. Values are normalized to the shared JSON wire contract at the event boundary. Unsupported short-term or sensory values are omitted from the generated event. Long-term observations are drained for the current partition key as one JSON batch; an invalid record is skipped after a valid batch is parsed, while a drain, serialization, or malformed-payload failure drops the batch. These observation failures do not change action success, and a long-term memory backend operation that already happened is not rolled back. + +The folding rules are: + +- one event per memory scope and operation kind; +- dot-separated memory paths for short-term and sensory `value` maps; +- last write or read wins when the same path is touched more than once by one action; +- operation-specific structured maps for long-term updates, gets, and searches; +- long-term search events keep the last result list for each memory-set/query pair. + +Actions triggered by memory events can use memory, but those memory operations do not generate another layer of memory events. This prevents memory-event subscriptions from recursively producing more memory events. + +## Configuration + +Memory event generation is controlled per operation kind. Each operation resolves its switch in this order: + +1. the operation-specific sub-key, if explicitly configured; +2. the `memory.generate-event` master switch, if explicitly configured; +3. the operation's built-in default. + +| Key | Effective default | Description | +|-----|-------------------|-------------| +| `memory.generate-event` | — | Master switch. Fallback for unset sub-keys; when unset itself, per-operation defaults apply. | +| `memory.generate-event.short-term-write` | on | Short-term memory writes. | +| `memory.generate-event.short-term-read` | off | Short-term memory reads. | +| `memory.generate-event.sensory-write` | on | Sensory memory writes. | +| `memory.generate-event.sensory-read` | off | Sensory memory reads. | +| `memory.generate-event.long-term-update` | on | Long-term memory adds, updates, and deletes. | +| `memory.generate-event.long-term-get` | on | Long-term memory gets. | +| `memory.generate-event.long-term-search` | on | Long-term memory searches. | +| `agent-run.begin-event` | false | Opt-in agent-run begin event. Independent of the `memory.generate-event` master switch. | + +The raw defaults of the eight `memory.generate-event*` options are unset; the on/off values above are the effective defaults after resolution. Registering an `EventListener` or subscribing an action to `_agent_run_begin_event` does not automatically enable the run-begin event. Set `agent-run.begin-event` to `true` when you need it. + +{{< tabs "Memory Event Configuration" >}} + +{{< tab "Python" >}} +```python +from flink_agents.api.core_options import AgentExecutionOptions, MemoryEventOptions + +config = agents_env.get_config() +# Turn everything off with the master switch... +config.set(MemoryEventOptions.MEMORY_GENERATE_EVENT, False) +# ...but keep short-term writes observable via the sub-key. +config.set(MemoryEventOptions.SHORT_TERM_WRITE, True) +# Opt in to the run-begin event through its independent switch. +config.set(AgentExecutionOptions.AGENT_RUN_BEGIN_EVENT, True) +``` +{{< /tab >}} + +{{< tab "Java" >}} +```java +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.configuration.MemoryEventOptions; + +Configuration config = agentsEnv.getConfig(); +// Turn everything off with the master switch... +config.set(MemoryEventOptions.MEMORY_GENERATE_EVENT, false); +// ...but keep short-term writes observable via the sub-key. +config.set(MemoryEventOptions.SHORT_TERM_WRITE, true); +// Opt in to the run-begin event through its independent switch. +config.set(AgentExecutionOptions.AGENT_RUN_BEGIN_EVENT, true); +``` +{{< /tab >}} + +{{< tab "YAML" >}} +```yaml +agent: + memory: + generate-event: false + generate-event.short-term-write: true + agent-run: + begin-event: true +``` +{{< /tab >}} + +{{< /tabs >}} + +## Event Types + +Flink Agents defines seven memory event types and one run-begin event: + +| Event type | Typed class | Emitted for | +|------------|-------------|-------------| +| `_short_term_write_event` | `ShortTermWriteEvent` | Short-term memory writes of one action | +| `_short_term_read_event` | `ShortTermReadEvent` | Short-term memory reads of one action | +| `_sensory_write_event` | `SensoryWriteEvent` | Sensory memory writes of one action | +| `_sensory_read_event` | `SensoryReadEvent` | Sensory memory reads of one action | +| `_long_term_update_event` | `LongTermUpdateEvent` | Long-term memory adds, updates, and deletes of one action | +| `_long_term_get_event` | `LongTermGetEvent` | Long-term memory gets of one action | +| `_long_term_search_event` | `LongTermSearchEvent` | Long-term memory searches of one action | +| `_agent_run_begin_event` | `AgentRunBeginEvent` | Begin of one agent run (STM value snapshot) | + +The seven memory event classes (Java: `org.apache.flink.agents.api.event`, Python: `flink_agents.api.events.memory_event`) share the abstract `MemoryEvent` base, and each class pins its own event type constant. `MemoryEvent.fromEvent(event)` (Python: `MemoryEvent.from_event(event)`) reconstructs the matching subclass from a generic event. + +`AgentRunBeginEvent` is a lifecycle event rather than a `MemoryEvent` subclass. It has the same `key` and `value` attribute shape, but it is controlled by the opt-in `agent-run.begin-event` switch instead of the memory-event switches. + +`LongTermSearchEvent` exposes the search shape as `getResults()` in Java and `.results` in Python: memory set to query string to ordered hit list. + +### Wire Format + +All memory events and run-begin events are emitted for keyed streams and carry two common attributes: + +| Attribute | Description | +|-----------|-------------| +| `key` | The key for this memory event, stored as text. Java keys and PyFlink keys with an explicit type normally use `String.valueOf`. PyFlink keys stored with pickle are decoded first and then converted with Python `str`. PyFlink byte-array keys with an explicit type use Python's bytes text directly and are not decoded with pickle. This value identifies the memory context. Flink does not use it to route records. | +| `value` | An operation-specific folded JSON map. Short-term, sensory, and run-begin snapshots use dot-key flat paths; long-term events use structured memory-set/item/query maps described below. | + +The `key` and `value` fields are nested under `attributes`; they are not top-level event fields: + +```json +{"timestamp": "2026-07-03T04:22:36.087914Z", "logLevel": "VERBOSE", "eventType": "_short_term_write_event", "event": {"eventType": "_short_term_write_event", "id": "da8ca017-25cc-4942-9085-ae38490e303c", "attributes": {"key": "user-42", "value": {"user.tier": "gold"}}, "type": "_short_term_write_event"}} +``` + +### Value Semantics + +- **Dot-key flat maps.** Nested short-term and sensory memory paths are flattened with dot-separated keys. Writing `user.tier` produces `{"user.tier": "gold"}`, not a nested object. +- **Net effect.** Multiple operations on the same field within one action fold to the last value. The event describes the action's net effect, not every intermediate step. +- **Short-term and sensory update records.** Write events reuse the updates recorded for memory persistence. `newObject(path)` and `set(path, null)` can therefore both appear as `path: null`; consumers cannot distinguish an object marker from a real null value from the event alone. +- **Long-term updates.** `_long_term_update_event` values use `{"": {"": value}}`. A `null` item value means the item was deleted. The separate `cleared_sets` attribute lists whole sets deleted during the action. A set clear removes earlier folded writes for that set; writes after the clear remain in `value`. +- **Long-term gets.** `_long_term_get_event` values use `{"": {"": value}}` for returned items. +- **Long-term searches.** `_long_term_search_event` uses `{"": {"": [{"id": ..., "value": ..., "score": ...}]}}`. If the same query is searched more than once in one set during one action, the last hit list is kept. +- **Read-only payloads.** Subscribers should treat `value` payloads as read-only. + +## Subscribing to Memory Events + +Memory events participate in normal event routing, so actions can subscribe to them by event type. + +{{< tabs "Subscribing to Memory Events" >}} + +{{< tab "Python" >}} +```python +from flink_agents.api.agents.agent import Agent +from flink_agents.api.decorators import action +from flink_agents.api.events.event import Event +from flink_agents.api.events.event_type import EventType +from flink_agents.api.events.memory_event import MemoryEvent +from flink_agents.api.runner_context import RunnerContext + + +class AuditingAgent(Agent): + @action(EventType.ShortTermWriteEvent) + @staticmethod + def on_stm_write(event: Event, ctx: RunnerContext) -> None: + write = MemoryEvent.from_event(event) + print(f"key={write.key} wrote {write.value}") +``` +{{< /tab >}} + +{{< tab "Java" >}} +```java +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.EventType; +import org.apache.flink.agents.api.annotation.Action; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.MemoryEvent; + +@Action(EventType.ShortTermWriteEvent) +public static void onStmWrite(Event event, RunnerContext ctx) { + MemoryEvent write = MemoryEvent.fromEvent(event); + System.out.printf("key=%s wrote %s%n", write.getKey(), write.getValue()); +} +``` +{{< /tab >}} + +{{< /tabs >}} + +## Inspecting Memory Values from the Event Log + +Short-term and sensory changes can be followed offline from the Event Log if the job logs events at `VERBOSE` level. Write events may contain ambiguous `null` object markers, while run-begin events contain only value nodes, so these events are not a complete reconstruction format. + +**Short-term memory** for a key at time *t*: + +1. Locate the latest `_agent_run_begin_event` for that key at or before *t*. Its `value` is the short-term value snapshot at run begin. +2. Apply subsequent `_short_term_write_event` values for the same key in log order until *t*. + +Because `agent-run.begin-event` is disabled by default, the Event Log has no built-in short-term value snapshot unless you opt in. + +**Sensory memory** starts empty at each run begin. Apply that run's `_sensory_write_event` values in log order until *t*. + +**Long-term memory cannot be reconstructed** from memory events. Long-term memory writes are model-mediated, so the requested update is not always the same as the final backend contents. Use long-term memory events to observe operation requests, trigger actions, and debug behavior, not as a backup of the long-term memory store. + +## Semantics & Caveats + +- **Event Log level.** Use `event-log.level: VERBOSE` when the full memory-event payload is required. +- **Durable execution.** A long-term memory operation wrapped in `durable_execute` is not recorded again when the durable result is replayed from state. The generated events describe operations that actually executed in the current attempt. +- **Observation failures.** Unsupported short-term or sensory values are omitted. Invalid long-term records are skipped individually after a valid batch is parsed, but a bridge, batch-serialization, or malformed-payload failure drops the current partition key's LTM observation batch. Warnings do not include observed values. These failures do not fail the action or roll back a completed long-term backend operation. +- **Failed actions.** A failed action does not reach the memory-event emission boundary. The operator task fails and recovery rebuilds the interpreter, so queued actions do not continue in the same interpreter instance after the failure. +- **Performance.** When enabled, the run-begin event scans the key's short-term memory on every input. Keep it disabled when you do not need offline value inspection. + +{{< hint warning >}} +**At-least-once on recovery.** After a failure, replayed completed actions can re-log their persisted output events, including generated memory events. A replayed input can also re-emit its run-begin event with restored short-term memory. Event Log consumers should tolerate duplicates. +{{< /hint >}} + +{{< hint warning >}} +**Short-term memory TTL.** With the default update type `ON_READ_AND_WRITE`, 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`. +{{< /hint >}} diff --git a/docs/content/docs/development/memory/sensory_and_short_term_memory.md b/docs/content/docs/development/memory/sensory_and_short_term_memory.md index a3b8b5e8cc..1109fb8742 100644 --- a/docs/content/docs/development/memory/sensory_and_short_term_memory.md +++ b/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 >}} +With the default update type `ON_READ_AND_WRITE`, 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 used by [Memory Events]({{< ref "docs/development/memory/memory_events" >}}), 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`. +{{< /hint >}} + {{< tabs "Short-Term Memory TTL Configuration" >}} {{< tab "Python" >}} diff --git a/docs/content/docs/operations/configuration.md b/docs/content/docs/operations/configuration.md index 9ba9e200ee..6622d94ac8 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -141,9 +141,25 @@ Here is the list of all built-in core configuration options. | `event-log.standard.max-array-elements` | 20 | int | At `STANDARD` level, arrays in the event payload with more than this many elements are truncated. Has no effect at `VERBOSE`. | | `event-log.standard.max-depth` | 5 | int | At `STANDARD` level, objects nested deeper than this are summarized. Has no effect at `VERBOSE`. | | `short-term-memory.state-ttl.ms` | 0 | long | Time-to-live for short-term memory state in milliseconds. Set to a value greater than 0 to enable TTL; 0 disables it. | -| `short-term-memory.state-ttl.update-type` | `ON_READ_AND_WRITE` | ShortTermMemoryTtlUpdate | Update policy for short-term memory TTL. Only applies when `short-term-memory.state-ttl.ms` is greater than 0. Valid values: `ON_CREATE_AND_WRITE`, `ON_READ_AND_WRITE`. | +| `short-term-memory.state-ttl.update-type` | `ON_READ_AND_WRITE` | ShortTermMemoryTtlUpdate | Update policy for short-term memory TTL. Only applies when `short-term-memory.state-ttl.ms` is greater than 0. Valid values: `ON_CREATE_AND_WRITE`, `ON_READ_AND_WRITE`. An enabled run-begin memory snapshot also refreshes TTL for entries it reads under `ON_READ_AND_WRITE`. | | `short-term-memory.state-ttl.visibility` | `NEVER_RETURN_EXPIRED` | ShortTermMemoryTtlVisibility | Visibility policy for expired short-term memory state. Only applies when `short-term-memory.state-ttl.ms` is greater than 0. Valid values: `NEVER_RETURN_EXPIRED`, `RETURN_EXPIRED_IF_NOT_CLEANED_UP`. | +### Memory Event Options + +The eight `memory.generate-event*` options have no raw `ConfigOption` default. When a sub-key and the master switch are both unset, the runtime uses the effective default shown below. See [Memory Events]({{< ref "docs/development/memory/memory_events" >}}) for resolution order, event payloads, and subscription examples. + +| Key | Raw default | Effective default | Type | Description | +|-----|-------------|-------------------|------|-------------| +| `memory.generate-event` | unset | per-operation defaults | boolean | Master fallback for unset operation-specific switches. | +| `memory.generate-event.short-term-write` | unset | on | boolean | Emit short-term memory write events. | +| `memory.generate-event.short-term-read` | unset | off | boolean | Emit short-term memory read events. | +| `memory.generate-event.sensory-write` | unset | on | boolean | Emit sensory memory write events. | +| `memory.generate-event.sensory-read` | unset | off | boolean | Emit sensory memory read events. | +| `memory.generate-event.long-term-update` | unset | on | boolean | Emit long-term memory add/delete events. | +| `memory.generate-event.long-term-get` | unset | on | boolean | Emit long-term memory get events. | +| `memory.generate-event.long-term-search` | unset | on | boolean | Emit long-term memory search events. | +| `agent-run.begin-event` | false | off | boolean | Opt in to the agent-run begin event. Independent of the memory-event master switch. | + ### Action State Store #### Common diff --git a/e2e-test/cross-language-event-snapshots/java/agent_run_begin_event.json b/e2e-test/cross-language-event-snapshots/java/agent_run_begin_event.json new file mode 100644 index 0000000000..5c050f0ac5 --- /dev/null +++ b/e2e-test/cross-language-event-snapshots/java/agent_run_begin_event.json @@ -0,0 +1,11 @@ +{ + "id" : "00000000-0000-0000-0000-000000000001", + "attributes" : { + "key" : "user-42", + "value" : { + "user.tier" : "gold", + "user.address.city" : "SF" + } + }, + "type" : "_agent_run_begin_event" +} diff --git a/e2e-test/cross-language-event-snapshots/java/short_term_write_event.json b/e2e-test/cross-language-event-snapshots/java/short_term_write_event.json new file mode 100644 index 0000000000..a6f906c2ff --- /dev/null +++ b/e2e-test/cross-language-event-snapshots/java/short_term_write_event.json @@ -0,0 +1,11 @@ +{ + "id" : "00000000-0000-0000-0000-000000000001", + "attributes" : { + "key" : "user-42", + "value" : { + "user.tier" : "gold", + "profile.m_a01" : null + } + }, + "type" : "_short_term_write_event" +} diff --git a/e2e-test/cross-language-event-snapshots/python/agent_run_begin_event.json b/e2e-test/cross-language-event-snapshots/python/agent_run_begin_event.json new file mode 100644 index 0000000000..79626d1d46 --- /dev/null +++ b/e2e-test/cross-language-event-snapshots/python/agent_run_begin_event.json @@ -0,0 +1,11 @@ +{ + "id": "00000000-0000-0000-0000-000000000001", + "type": "_agent_run_begin_event", + "attributes": { + "key": "user-42", + "value": { + "user.tier": "gold", + "user.address.city": "SF" + } + } +} diff --git a/e2e-test/cross-language-event-snapshots/python/short_term_write_event.json b/e2e-test/cross-language-event-snapshots/python/short_term_write_event.json new file mode 100644 index 0000000000..1a4ac90a94 --- /dev/null +++ b/e2e-test/cross-language-event-snapshots/python/short_term_write_event.json @@ -0,0 +1,11 @@ +{ + "id": "00000000-0000-0000-0000-000000000001", + "type": "_short_term_write_event", + "attributes": { + "key": "user-42", + "value": { + "user.tier": "gold", + "profile.m_a01": null + } + } +} diff --git a/python/flink_agents/api/core_options.py b/python/flink_agents/api/core_options.py index d560d8cf93..d5cd70da6c 100644 --- a/python/flink_agents/api/core_options.py +++ b/python/flink_agents/api/core_options.py @@ -231,6 +231,39 @@ class AgentConfigOptions: ) +class MemoryEventOptions: + """Options controlling memory observation events. + + Per-op resolution: sub-key explicit -> master switch explicit -> built-in default + (writes and long-term ops on; short-term/sensory reads off) + """ + + MEMORY_GENERATE_EVENT = ConfigOption( + key="memory.generate-event", config_type=bool, default=None + ) + SHORT_TERM_WRITE = ConfigOption( + key="memory.generate-event.short-term-write", config_type=bool, default=None + ) + SHORT_TERM_READ = ConfigOption( + key="memory.generate-event.short-term-read", config_type=bool, default=None + ) + SENSORY_WRITE = ConfigOption( + key="memory.generate-event.sensory-write", config_type=bool, default=None + ) + SENSORY_READ = ConfigOption( + key="memory.generate-event.sensory-read", config_type=bool, default=None + ) + LONG_TERM_UPDATE = ConfigOption( + key="memory.generate-event.long-term-update", config_type=bool, default=None + ) + LONG_TERM_GET = ConfigOption( + key="memory.generate-event.long-term-get", config_type=bool, default=None + ) + LONG_TERM_SEARCH = ConfigOption( + key="memory.generate-event.long-term-search", config_type=bool, default=None + ) + + class AgentExecutionOptions: """Execution options for Flink Agents.""" @@ -264,6 +297,10 @@ class AgentExecutionOptions: default=True, ) + AGENT_RUN_BEGIN_EVENT = ConfigOption( + key="agent-run.begin-event", config_type=bool, default=False + ) + TOOL_CALL_ASYNC = ConfigOption( key="tool-call.async", config_type=bool, diff --git a/python/flink_agents/api/events/event_type.py b/python/flink_agents/api/events/event_type.py index a244661a77..869fe12f37 100644 --- a/python/flink_agents/api/events/event_type.py +++ b/python/flink_agents/api/events/event_type.py @@ -41,6 +41,30 @@ from flink_agents.api.events.event import ( OutputEvent as _OutputEvent, ) +from flink_agents.api.events.memory_event import ( + LongTermGetEvent as _LongTermGetEvent, +) +from flink_agents.api.events.memory_event import ( + LongTermSearchEvent as _LongTermSearchEvent, +) +from flink_agents.api.events.memory_event import ( + LongTermUpdateEvent as _LongTermUpdateEvent, +) +from flink_agents.api.events.memory_event import ( + SensoryReadEvent as _SensoryReadEvent, +) +from flink_agents.api.events.memory_event import ( + SensoryWriteEvent as _SensoryWriteEvent, +) +from flink_agents.api.events.memory_event import ( + ShortTermReadEvent as _ShortTermReadEvent, +) +from flink_agents.api.events.memory_event import ( + ShortTermWriteEvent as _ShortTermWriteEvent, +) +from flink_agents.api.events.run_event import ( + AgentRunBeginEvent as _AgentRunBeginEvent, +) from flink_agents.api.events.tool_event import ( ToolRequestEvent as _ToolRequestEvent, ) @@ -63,3 +87,11 @@ class EventType: ToolResponseEvent: str = _ToolResponseEvent.EVENT_TYPE ContextRetrievalRequestEvent: str = _ContextRetrievalRequestEvent.EVENT_TYPE ContextRetrievalResponseEvent: str = _ContextRetrievalResponseEvent.EVENT_TYPE + ShortTermWriteEvent: str = _ShortTermWriteEvent.EVENT_TYPE + ShortTermReadEvent: str = _ShortTermReadEvent.EVENT_TYPE + SensoryWriteEvent: str = _SensoryWriteEvent.EVENT_TYPE + SensoryReadEvent: str = _SensoryReadEvent.EVENT_TYPE + LongTermUpdateEvent: str = _LongTermUpdateEvent.EVENT_TYPE + LongTermGetEvent: str = _LongTermGetEvent.EVENT_TYPE + LongTermSearchEvent: str = _LongTermSearchEvent.EVENT_TYPE + AgentRunBeginEvent: str = _AgentRunBeginEvent.EVENT_TYPE diff --git a/python/flink_agents/api/events/memory_event.py b/python/flink_agents/api/events/memory_event.py new file mode 100644 index 0000000000..6be5a64a44 --- /dev/null +++ b/python/flink_agents/api/events/memory_event.py @@ -0,0 +1,204 @@ +################################################################################ +# 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. +################################################################################# +"""Memory observation events, mirroring the Java MemoryEvent hierarchy.""" + +import base64 +import json +from typing import Any, ClassVar, Dict, List + +from typing_extensions import override + +from flink_agents.api.events.event import Event + + +def _observation_attributes(key: Any, value: Any) -> Dict[str, Any]: + """Build attributes shared by memory-operation and run-begin events.""" + if not isinstance(key, str): + msg = "Memory observation attribute 'key' must be a string" + raise TypeError(msg) + if not isinstance(value, dict): + msg = "Memory observation attribute 'value' must be a dict" + raise TypeError(msg) + try: + normalized = json.loads( + json.dumps( + value, + ensure_ascii=False, + allow_nan=False, + default=_encode_json_value, + ) + ) + except (TypeError, ValueError) as error: + msg = "Memory observation value must be JSON serializable" + raise ValueError(msg) from error + return {"key": key, "value": normalized} + + +def _encode_json_value(value: Any) -> Any: + """Mirror Jackson's byte[] wire representation; reject other unknown values.""" + if isinstance(value, bytes | bytearray): + return base64.b64encode(bytes(value)).decode("ascii") + msg = f"Object of type {type(value).__name__} is not JSON serializable" + raise TypeError(msg) + + +class MemoryEvent(Event): + """Base class of the memory observation events, emitted at the action + finish boundary. Each concrete subclass pins one of the operation kinds + as its ``EVENT_TYPE``; instantiate a subclass, not this base class. + + Attributes ``key`` (the String Flink key) and ``value`` (the operation's + folded JSON value map) live inside ``attributes``: + ``{id, type, attributes: {key, value}}``. + """ + + EVENT_TYPE: ClassVar[str | None] = None # pinned by each concrete subclass + + def __init__(self, *, key: str, value: Dict[str, Any]) -> None: + """Create a memory event; the type comes from the subclass EVENT_TYPE.""" + event_type = type(self).EVENT_TYPE + if event_type is None: + msg = "MemoryEvent is abstract; instantiate one of its subclasses" + raise TypeError(msg) + super().__init__( + type=event_type, attributes=_observation_attributes(key, value) + ) + + @classmethod + def is_memory_type(cls, type: str | None) -> bool: + """Return True iff type is one of the seven memory operation types.""" + return type in _TYPE_TO_CLASS + + @classmethod + @override + def from_event(cls, event: Event) -> "MemoryEvent": + """Reconstruct the typed subclass view from a base Event.""" + subclass = _TYPE_TO_CLASS.get(event.type) + if subclass is None: + msg = f"Not a memory event type: {event.type}" + raise ValueError(msg) + kwargs = { + "key": event.attributes.get("key"), + "value": event.attributes.get("value"), + } + if subclass is LongTermUpdateEvent: + kwargs["cleared_sets"] = event.attributes.get("cleared_sets", []) + result = subclass(**kwargs) + return result.reconstruct_from(event) + + @property + def key(self) -> str: + """Return the String Flink key this operation belongs to.""" + return self.get_attr("key") + + @property + def value(self) -> Dict[str, Any]: + """Return this operation's folded JSON map.""" + return self.get_attr("value") + + +class ShortTermWriteEvent(MemoryEvent): + """Memory observation event: a write to short-term memory.""" + + EVENT_TYPE: ClassVar[str] = "_short_term_write_event" + + +class ShortTermReadEvent(MemoryEvent): + """Memory observation event: a read from short-term memory.""" + + EVENT_TYPE: ClassVar[str] = "_short_term_read_event" + + +class SensoryWriteEvent(MemoryEvent): + """Memory observation event: a write to sensory memory.""" + + EVENT_TYPE: ClassVar[str] = "_sensory_write_event" + + +class SensoryReadEvent(MemoryEvent): + """Memory observation event: a read from sensory memory.""" + + EVENT_TYPE: ClassVar[str] = "_sensory_read_event" + + +class LongTermUpdateEvent(MemoryEvent): + """Long-term memory adds, updates, item deletes, and set clears.""" + + EVENT_TYPE: ClassVar[str] = "_long_term_update_event" + + def __init__( + self, + *, + key: str, + value: Dict[str, Any], + cleared_sets: List[str] | None = None, + ) -> None: + """Create an update event with item changes and explicit set clears.""" + if cleared_sets is not None and not isinstance(cleared_sets, list): + msg = "Long-term update 'cleared_sets' must be a list" + raise TypeError(msg) + if cleared_sets is not None and not all( + isinstance(memory_set, str) for memory_set in cleared_sets + ): + msg = "Long-term update 'cleared_sets' must contain only strings" + raise TypeError(msg) + super().__init__(key=key, value=value) + self.attributes = { + **self.attributes, + "cleared_sets": list(cleared_sets) if cleared_sets is not None else [], + } + + @property + def cleared_sets(self) -> List[str]: + """Memory sets cleared before any later folded writes in this event.""" + return self.get_attr("cleared_sets") + + +class LongTermGetEvent(MemoryEvent): + """Memory observation event: a long-term memory read by id.""" + + EVENT_TYPE: ClassVar[str] = "_long_term_get_event" + + +class LongTermSearchEvent(MemoryEvent): + """Memory observation event: a semantic search over long-term memory. + + ``value`` maps each memory set to a map from query string to ordered hit + list; ``results`` exposes that shape. + """ + + EVENT_TYPE: ClassVar[str] = "_long_term_search_event" + + @property + def results(self) -> Dict[str, Dict[str, List[Dict[str, Any]]]]: + """Typed view: memory set → query → ordered hit list.""" + return self.get_attr("value") + + +_TYPE_TO_CLASS: Dict[str, type[MemoryEvent]] = { + subclass.EVENT_TYPE: subclass + for subclass in ( + ShortTermWriteEvent, + ShortTermReadEvent, + SensoryWriteEvent, + SensoryReadEvent, + LongTermUpdateEvent, + LongTermGetEvent, + LongTermSearchEvent, + ) +} diff --git a/python/flink_agents/api/events/run_event.py b/python/flink_agents/api/events/run_event.py new file mode 100644 index 0000000000..b8948c22e8 --- /dev/null +++ b/python/flink_agents/api/events/run_event.py @@ -0,0 +1,61 @@ +################################################################################ +# 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. +################################################################################# +"""Agent-run lifecycle events, mirroring the Java classes.""" + +from typing import Any, ClassVar, Dict + +from typing_extensions import override + +from flink_agents.api.events.event import Event +from flink_agents.api.events.memory_event import _observation_attributes + + +class AgentRunBeginEvent(Event): + """Lifecycle event carrying one run's STM value snapshot. + + When ``agent-run.begin-event`` is enabled, the framework emits this event + before the run's actions execute. + """ + + EVENT_TYPE: ClassVar[str] = "_agent_run_begin_event" + + def __init__(self, *, key: str, value: Dict[str, Any]) -> None: + """Create an AgentRunBeginEvent carrying STM VALUE nodes for the key.""" + super().__init__( + type=AgentRunBeginEvent.EVENT_TYPE, + attributes=_observation_attributes(key, value), + ) + + @classmethod + @override + def from_event(cls, event: Event) -> "AgentRunBeginEvent": + """Rebuild the typed view from a generic Event of this type.""" + result = AgentRunBeginEvent( + key=event.attributes.get("key"), value=event.attributes.get("value") + ) + return result.reconstruct_from(event) + + @property + def key(self) -> str: + """Return the String Flink key of this run.""" + return self.get_attr("key") + + @property + def value(self) -> Dict[str, Any]: + """Return the STM value snapshot as a dot-key flat map.""" + return self.get_attr("value") diff --git a/python/flink_agents/api/tests/test_cross_language_event_snapshots.py b/python/flink_agents/api/tests/test_cross_language_event_snapshots.py index 83648967a2..545cbfdf49 100644 --- a/python/flink_agents/api/tests/test_cross_language_event_snapshots.py +++ b/python/flink_agents/api/tests/test_cross_language_event_snapshots.py @@ -32,6 +32,8 @@ ContextRetrievalResponseEvent, ) from flink_agents.api.events.event import Event, InputEvent, OutputEvent +from flink_agents.api.events.memory_event import MemoryEvent, ShortTermWriteEvent +from flink_agents.api.events.run_event import AgentRunBeginEvent from flink_agents.api.events.tool_event import ToolRequestEvent, ToolResponseEvent from flink_agents.api.vector_stores.vector_store import Document @@ -164,7 +166,9 @@ def test_python_can_deserialize_chat_request_event_from_java_snapshot() -> None: assert msg.content == "hello world" -def test_chat_request_row_type_info_output_schema_is_not_portable_across_languages_known_gap() -> None: +def test_chat_request_row_type_info_output_schema_is_not_portable_across_languages_known_gap() -> ( + None +): """Known 0.3 gap — RowTypeInfo-typed output_schema does not round-trip across the language boundary. Python emits ``{"names": [...], "types": []}`` while Java emits ``{"fieldNames": [...], "types": []}``, so a ChatRequestEvent carrying a @@ -190,8 +194,8 @@ def test_chat_request_row_type_info_output_schema_is_not_portable_across_languag payload = event.model_dump_json() # Pin Python's local shape so a future regression can't silently change it. The gap with # Java's `{"fieldNames": ...}` shape is the documented limitation, not the assertion. - assert "\"names\"" in payload - assert "\"fieldNames\"" not in payload + assert '"names"' in payload + assert '"fieldNames"' not in payload # ── ChatResponseEvent ─────────────────────────────────────────────────── @@ -222,7 +226,9 @@ def test_python_can_deserialize_chat_response_event_from_java_snapshot() -> None typed = ChatResponseEvent.from_event(base) expected_request_id = str(_FIXED_REQUEST_ID) actual_request_id = ( - str(typed.request_id) if not isinstance(typed.request_id, str) else typed.request_id + str(typed.request_id) + if not isinstance(typed.request_id, str) + else typed.request_id ) assert actual_request_id == expected_request_id, "request_id mismatch." assert typed.response is not None, "response is None." @@ -236,7 +242,11 @@ def test_python_can_deserialize_chat_response_event_from_java_snapshot() -> None def _build_tool_request_event() -> ToolRequestEvent: - tool_call = {"id": _FIXED_TOOL_CALL_ID, "name": "echo", "arguments": {"value": "ping"}} + tool_call = { + "id": _FIXED_TOOL_CALL_ID, + "name": "echo", + "arguments": {"value": "ping"}, + } event = ToolRequestEvent(model="test-model", tool_calls=[tool_call]) return _force_id(event, _FIXED_EVENT_ID) @@ -339,7 +349,9 @@ def test_context_retrieval_request_event_python_snapshot_is_stable() -> None: ) -def test_python_can_deserialize_context_retrieval_request_event_from_java_snapshot() -> None: +def test_python_can_deserialize_context_retrieval_request_event_from_java_snapshot() -> ( + None +): base = _read_java_snapshot("context_retrieval_request_event.json") typed = ContextRetrievalRequestEvent.from_event(base) assert typed.query == "what is flink" @@ -376,12 +388,16 @@ def test_context_retrieval_response_event_python_snapshot_is_stable() -> None: ) -def test_python_can_deserialize_context_retrieval_response_event_from_java_snapshot() -> None: +def test_python_can_deserialize_context_retrieval_response_event_from_java_snapshot() -> ( + None +): base = _read_java_snapshot("context_retrieval_response_event.json") typed = ContextRetrievalResponseEvent.from_event(base) expected_request_id = str(_FIXED_REQUEST_ID) actual_request_id = ( - str(typed.request_id) if not isinstance(typed.request_id, str) else typed.request_id + str(typed.request_id) + if not isinstance(typed.request_id, str) + else typed.request_id ) assert actual_request_id == expected_request_id assert typed.query == "what is flink" @@ -390,6 +406,77 @@ def test_python_can_deserialize_context_retrieval_response_event_from_java_snaps assert typed.documents[0].id == "doc-1" +# ── Memory observation events ────────────────────────────────────────── + + +def _build_short_term_write_event() -> ShortTermWriteEvent: + return _force_id( + ShortTermWriteEvent( + key="user-42", + value={"user.tier": "gold", "profile.m_a01": None}, + ), + _FIXED_EVENT_ID, + ) + + +def test_regenerate_short_term_write_event_python_snapshot() -> None: + if not _regenerate_enabled(): + pytest.skip("Set REGENERATE_SNAPSHOTS=1 to refresh.") + _write_python_snapshot( + "short_term_write_event.json", _build_short_term_write_event() + ) + + +def test_short_term_write_event_python_snapshot_is_stable() -> None: + _assert_python_snapshot_stable( + "short_term_write_event.json", _build_short_term_write_event() + ) + + +def test_python_can_deserialize_short_term_write_event_from_java_snapshot() -> None: + base = _read_java_snapshot("short_term_write_event.json") + typed = MemoryEvent.from_event(base) + + assert isinstance(typed, ShortTermWriteEvent) + assert typed.id == _FIXED_EVENT_ID + assert typed.key == "user-42" + assert typed.value == {"user.tier": "gold", "profile.m_a01": None} + + +# ── Agent-run lifecycle events ───────────────────────────────────────── + + +def _build_agent_run_begin_event() -> AgentRunBeginEvent: + return _force_id( + AgentRunBeginEvent( + key="user-42", + value={"user.tier": "gold", "user.address.city": "SF"}, + ), + _FIXED_EVENT_ID, + ) + + +def test_regenerate_agent_run_begin_event_python_snapshot() -> None: + if not _regenerate_enabled(): + pytest.skip("Set REGENERATE_SNAPSHOTS=1 to refresh.") + _write_python_snapshot("agent_run_begin_event.json", _build_agent_run_begin_event()) + + +def test_agent_run_begin_event_python_snapshot_is_stable() -> None: + _assert_python_snapshot_stable( + "agent_run_begin_event.json", _build_agent_run_begin_event() + ) + + +def test_python_can_deserialize_agent_run_begin_event_from_java_snapshot() -> None: + base = _read_java_snapshot("agent_run_begin_event.json") + typed = AgentRunBeginEvent.from_event(base) + + assert typed.id == _FIXED_EVENT_ID + assert typed.key == "user-42" + assert typed.value == {"user.tier": "gold", "user.address.city": "SF"} + + # ── Generic Event with primitive attributes (user-authored axis) ─────── diff --git a/python/flink_agents/api/tests/test_memory_event.py b/python/flink_agents/api/tests/test_memory_event.py new file mode 100644 index 0000000000..211be33ada --- /dev/null +++ b/python/flink_agents/api/tests/test_memory_event.py @@ -0,0 +1,230 @@ +################################################################################ +# 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. +################################################################################# +import pytest + +from flink_agents.api.configuration import ConfigOption +from flink_agents.api.core_options import AgentExecutionOptions, MemoryEventOptions +from flink_agents.api.events.event import Event +from flink_agents.api.events.event_type import EventType +from flink_agents.api.events.memory_event import ( + LongTermGetEvent, + LongTermSearchEvent, + LongTermUpdateEvent, + MemoryEvent, + SensoryReadEvent, + SensoryWriteEvent, + ShortTermReadEvent, + ShortTermWriteEvent, +) + + +def test_subclasses_pin_types_match_java() -> None: + assert ShortTermWriteEvent.EVENT_TYPE == "_short_term_write_event" + assert ShortTermReadEvent.EVENT_TYPE == "_short_term_read_event" + assert SensoryWriteEvent.EVENT_TYPE == "_sensory_write_event" + assert SensoryReadEvent.EVENT_TYPE == "_sensory_read_event" + assert LongTermUpdateEvent.EVENT_TYPE == "_long_term_update_event" + assert LongTermGetEvent.EVENT_TYPE == "_long_term_get_event" + assert LongTermSearchEvent.EVENT_TYPE == "_long_term_search_event" + + assert ShortTermWriteEvent(key="k", value={}).type == "_short_term_write_event" + assert LongTermSearchEvent(key="k", value={}).type == "_long_term_search_event" + + +def test_is_memory_type() -> None: + assert MemoryEvent.is_memory_type("_short_term_write_event") + assert not MemoryEvent.is_memory_type("_input_event") + assert not MemoryEvent.is_memory_type(None) + + +def test_event_type_aggregate() -> None: + assert EventType.ShortTermWriteEvent == ShortTermWriteEvent.EVENT_TYPE + assert EventType.LongTermSearchEvent == LongTermSearchEvent.EVENT_TYPE + + +def test_key_value_in_attributes() -> None: + e = ShortTermWriteEvent(key="user-42", value={"user.tier": "gold"}) + assert e.attributes["key"] == "user-42" + assert e.attributes["value"] == {"user.tier": "gold"} + assert e.key == "user-42" + assert e.value["user.tier"] == "gold" + + +def test_values_are_normalized_and_deep_copied() -> None: + original = {"nested": {"value": 1}, "bytes": b"\x01\x02\x03"} + event = ShortTermWriteEvent(key="k", value=original) + + original["nested"]["value"] = 2 + assert event.value == {"nested": {"value": 1}, "bytes": "AQID"} + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_non_finite_values_are_rejected(value: float) -> None: + with pytest.raises(ValueError, match="JSON serializable"): + ShortTermWriteEvent(key="k", value={"bad": value}) + + +def test_base_class_is_abstract() -> None: + with pytest.raises(TypeError, match="abstract"): + MemoryEvent(key="k", value={}) + + +def test_from_event_rejects_non_memory_type() -> None: + generic = Event(type="_input_event", attributes={"key": "k", "value": {}}) + with pytest.raises(ValueError, match="Not a memory event type"): + MemoryEvent.from_event(generic) + + +def test_from_event_dispatches_to_subclass() -> None: + generic = Event( + type=LongTermGetEvent.EVENT_TYPE, + attributes={"key": "k", "value": {"s.m1": "v"}}, + ) + typed = MemoryEvent.from_event(generic) + assert isinstance(typed, LongTermGetEvent) + assert typed.id == generic.id + assert typed.value == {"s.m1": "v"} + + +def test_constructors_reject_incomplete_attributes() -> None: + with pytest.raises(TypeError, match="key"): + ShortTermWriteEvent(key=None, value={}) # type: ignore[arg-type] + with pytest.raises(TypeError, match="value"): + ShortTermWriteEvent(key="k", value=None) # type: ignore[arg-type] + + generic = Event(type=ShortTermWriteEvent.EVENT_TYPE, attributes={"key": "k"}) + with pytest.raises(TypeError, match="value"): + MemoryEvent.from_event(generic) + + +def test_long_term_search_typed_results() -> None: + e = LongTermSearchEvent( + key="user-42", + value={ + "policies": { + "refund policy": [ + {"id": "p_01", "value": "7-day refund", "score": 0.92}, + { + "id": "p_02", + "value": "no custom refunds", + "score": 0.81, + }, + ] + } + }, + ) + assert len(e.results["policies"]["refund policy"]) == 2 + assert e.results["policies"]["refund policy"][0]["id"] == "p_01" + + +def test_long_term_update_carries_explicit_cleared_sets() -> None: + event = LongTermUpdateEvent( + key="k", value={"prefs": {"m2": "new"}}, cleared_sets=["prefs"] + ) + restored = MemoryEvent.from_event(event) + assert isinstance(restored, LongTermUpdateEvent) + assert restored.cleared_sets == ["prefs"] + + +@pytest.mark.parametrize( + "invalid_cleared_sets", + ["prefs", ("prefs",), {"prefs"}, {"prefs": True}, 1], +) +def test_long_term_update_requires_cleared_sets_list( + invalid_cleared_sets: object, +) -> None: + with pytest.raises(TypeError, match="must be a list"): + LongTermUpdateEvent( + key="k", + value={}, + cleared_sets=invalid_cleared_sets, # type: ignore[arg-type] + ) + + +def test_long_term_update_allows_none_but_rejects_non_string_list_items() -> None: + assert LongTermUpdateEvent(key="k", value={}, cleared_sets=None).cleared_sets == [] + with pytest.raises(TypeError, match="only strings"): + LongTermUpdateEvent( + key="k", + value={}, + cleared_sets=["prefs", 1], # type: ignore[list-item] + ) + + +def test_from_event_rejects_string_cleared_sets() -> None: + generic = Event( + type=LongTermUpdateEvent.EVENT_TYPE, + attributes={"key": "k", "value": {}, "cleared_sets": "prefs"}, + ) + + with pytest.raises(TypeError, match="must be a list"): + MemoryEvent.from_event(generic) + + +def test_config_option_keys_match_java() -> None: + # (option, expected key, expected default) for all memory options. These keys + # and defaults are the cross-language contract shared with Java + # MemoryEventOptions; every one must be asserted so a Python-side typo or + # drift on any sub-key is caught, not just the master switch. + expected = [ + (MemoryEventOptions.MEMORY_GENERATE_EVENT, "memory.generate-event", None), + ( + MemoryEventOptions.SHORT_TERM_WRITE, + "memory.generate-event.short-term-write", + None, + ), + ( + MemoryEventOptions.SHORT_TERM_READ, + "memory.generate-event.short-term-read", + None, + ), + ( + MemoryEventOptions.SENSORY_WRITE, + "memory.generate-event.sensory-write", + None, + ), + ( + MemoryEventOptions.SENSORY_READ, + "memory.generate-event.sensory-read", + None, + ), + ( + MemoryEventOptions.LONG_TERM_UPDATE, + "memory.generate-event.long-term-update", + None, + ), + ( + MemoryEventOptions.LONG_TERM_GET, + "memory.generate-event.long-term-get", + None, + ), + ( + MemoryEventOptions.LONG_TERM_SEARCH, + "memory.generate-event.long-term-search", + None, + ), + ] + for option, key, default in expected: + assert isinstance(option, ConfigOption) + assert option.get_key() == key + assert option.get_default_value() is default + + assert AgentExecutionOptions.AGENT_RUN_BEGIN_EVENT.get_key() == ( + "agent-run.begin-event" + ) + assert AgentExecutionOptions.AGENT_RUN_BEGIN_EVENT.get_default_value() is False diff --git a/python/flink_agents/api/tests/test_run_event.py b/python/flink_agents/api/tests/test_run_event.py new file mode 100644 index 0000000000..1dadeaeea2 --- /dev/null +++ b/python/flink_agents/api/tests/test_run_event.py @@ -0,0 +1,51 @@ +################################################################################ +# 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. +################################################################################# +import pytest + +from flink_agents.api.events.event import Event +from flink_agents.api.events.event_type import EventType +from flink_agents.api.events.run_event import AgentRunBeginEvent + + +def test_type_constant() -> None: + assert AgentRunBeginEvent.EVENT_TYPE == "_agent_run_begin_event" + assert EventType.AgentRunBeginEvent == AgentRunBeginEvent.EVENT_TYPE + + +def test_key_value() -> None: + e = AgentRunBeginEvent(key="user-42", value={"user.tier": "gold"}) + assert e.key == "user-42" + assert e.value == {"user.tier": "gold"} + assert e.attributes["key"] == "user-42" + + +def test_from_event() -> None: + generic = Event( + type=AgentRunBeginEvent.EVENT_TYPE, + attributes={"key": "k", "value": {"a": 1}}, + ) + typed = AgentRunBeginEvent.from_event(generic) + assert typed.id == generic.id + assert typed.value == {"a": 1} + + +def test_from_event_missing_attributes_raises() -> None: + # Missing "value" -> from_event must reject rather than silently build a bad event. + generic = Event(type=AgentRunBeginEvent.EVENT_TYPE, attributes={"key": "k"}) + with pytest.raises(TypeError): + AgentRunBeginEvent.from_event(generic) diff --git a/python/flink_agents/cli/tests/test_trace_tree.py b/python/flink_agents/cli/tests/test_trace_tree.py index ffcf0ce6e8..754896a570 100644 --- a/python/flink_agents/cli/tests/test_trace_tree.py +++ b/python/flink_agents/cli/tests/test_trace_tree.py @@ -114,6 +114,36 @@ def test_reader_reconstructs_text_and_json_in_log_order(tmp_path: Path) -> None: assert text_result.stderr == "" +def test_reader_keeps_enabled_run_begin_branch_in_input_trace(tmp_path: Path) -> None: + log_path = tmp_path / "events.log" + _write_log( + log_path, + [ + _record("root", "_input_event"), + _record( + "run-begin", + "_agent_run_begin_event", + "root", + "agent_run_begin_action", + ), + _record("output", "_output_event", "run-begin", "on_run_begin"), + ], + ) + + result = _run_reader(log_path, "json") + trace_forest = json.loads(result.stdout) + + assert trace_forest["roots"] == ["root"] + assert trace_forest["nodes"]["root"]["actions"] == [ + {"name": "agent_run_begin_action", "children": ["run-begin"]} + ] + assert trace_forest["nodes"]["run-begin"]["actions"] == [ + {"name": "on_run_begin", "children": ["output"]} + ] + assert trace_forest["warnings"] == [] + assert result.stderr == "" + + def test_reader_reconstructs_pretty_printed_records(tmp_path: Path) -> None: log_path = tmp_path / "events.log" _write_pretty_log( diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/memory_event_logging_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/memory_event_logging_test.py new file mode 100644 index 0000000000..0e515f1da3 --- /dev/null +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/memory_event_logging_test.py @@ -0,0 +1,118 @@ +################################################################################ +# 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. +################################################################################# +"""E2e test: memory events and agent-run-begin events land in the EventLog.""" + +import json +import os +import sysconfig +from pathlib import Path + +from pyflink.common import Configuration +from pyflink.datastream import StreamExecutionEnvironment + +from flink_agents.api.agents.agent import Agent +from flink_agents.api.core_options import AgentExecutionOptions +from flink_agents.api.decorators import action +from flink_agents.api.events.event import Event, InputEvent, OutputEvent +from flink_agents.api.events.event_type import EventType +from flink_agents.api.execution_environment import AgentsExecutionEnvironment +from flink_agents.api.runner_context import RunnerContext + +# Include both the active environment and any source checkout already on PYTHONPATH so +# the embedded interpreter can resolve this test module. +_purelib = sysconfig.get_paths()["purelib"] +_extra_pythonpath = os.environ.get("PYTHONPATH") +os.environ["PYTHONPATH"] = ( + f"{_purelib}{os.pathsep}{_extra_pythonpath}" if _extra_pythonpath else _purelib +) + +class MemoryWritingAgent(Agent): + """Agent whose input action writes short-term memory.""" + + @action(EventType.InputEvent) + @staticmethod + def process_input(event: Event, ctx: RunnerContext) -> None: + """Write STM and send an output event.""" + input_data = InputEvent.from_event(event).input + ctx.short_term_memory.set("user.tier", "gold") + ctx.send_event(OutputEvent(output={"processed": input_data["review"]})) + + +def _read_ordered_records(event_log_dir: Path) -> list[dict]: + """Read JSON records from the single event log, preserving line order.""" + log_files = list(event_log_dir.glob("events-*.log")) + assert len(log_files) == 1, ( + f"Expected one event log file in {event_log_dir}, found {len(log_files)}" + ) + with log_files[0].open(encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + + +def test_memory_event_logging(tmp_path: Path) -> None: + """Default memory events plus an opt-in run-begin event land in the Event Log.""" + event_log_dir = tmp_path / "event_log" + + # All inputs share one key so run-begin STM snapshots accumulate deterministically. + inputs = [{"id": 7, "review": f"input-{i}"} for i in range(3)] + + config = Configuration() + env = StreamExecutionEnvironment.get_execution_environment(config) + env.set_parallelism(1) + + agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env) + agents_env.get_config().set_str("baseLogDir", str(event_log_dir)) + # Memory operation events stay at defaults; run-begin snapshots are opt-in. + agents_env.get_config().set(AgentExecutionOptions.AGENT_RUN_BEGIN_EVENT, True) + # VERBOSE keeps the full event subtree. + agents_env.get_config().set_str("event-log.level", "VERBOSE") + + input_datastream = env.from_collection(inputs) + output_datastream = ( + agents_env.from_datastream( + input=input_datastream, key_selector=lambda value: value["id"] + ) + .apply(MemoryWritingAgent()) + .to_datastream() + ) + list(output_datastream.execute_and_collect()) + + records = _read_ordered_records(event_log_dir) + + types = [r["eventType"] for r in records] + assert EventType.AgentRunBeginEvent in types + assert EventType.ShortTermWriteEvent in types + + write = next(r for r in records if r["eventType"] == EventType.ShortTermWriteEvent) + assert write["event"]["attributes"]["value"]["user.tier"] == "gold" + assert write["event"]["attributes"]["key"] == "7" + + begins = [r for r in records if r["eventType"] == EventType.AgentRunBeginEvent] + assert len(begins) == len(inputs) + assert {event["event"]["attributes"]["key"] for event in begins} == {"7"} + # First run-begin: empty STM; a later one contains the previous run's write. + assert begins[0]["event"]["attributes"]["value"] == {} + assert begins[-1]["event"]["attributes"]["value"].get("user.tier") == "gold" + + # Adjacency: each run-begin line directly follows its InputEvent line (both are + # emitted synchronously back-to-back) — the reconstruction anchor. + for i, record in enumerate(records): + if record["eventType"] == EventType.AgentRunBeginEvent: + assert i > 0, "run-begin must not be the first log line" + assert records[i - 1]["eventType"] == EventType.InputEvent, ( + f"run-begin at line {i} should directly follow an _input_event line" + ) diff --git a/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py b/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py index 0e2762d73c..a1750a3f62 100644 --- a/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py +++ b/python/flink_agents/plan/tests/compatibility/check_java_python_config_options_parity.py @@ -32,7 +32,11 @@ from pyflink.util.java_utils import add_jars_to_context_class_loader from flink_agents.api.configuration import ConfigOption -from flink_agents.api.core_options import AgentConfigOptions, AgentExecutionOptions +from flink_agents.api.core_options import ( + AgentConfigOptions, + AgentExecutionOptions, + MemoryEventOptions, +) _JAVA_PRIMITIVE_TYPE_TO_PYTHON: dict[str, type] = { "java.lang.String": str, @@ -45,7 +49,9 @@ } -def collect_python_config_options(python_options_class: type) -> dict[str, ConfigOption]: +def collect_python_config_options( + python_options_class: type, +) -> dict[str, ConfigOption]: """Return ``{FIELD_NAME: ConfigOption}`` declared on a Python options class.""" return { name: value @@ -91,7 +97,9 @@ def normalize_java_default(java_default: Any, python_config_type: type) -> Any: if hasattr(java_default, "name") and callable(java_default.name): enum_name = java_default.name() - if isinstance(python_config_type, type) and issubclass(python_config_type, Enum): + if isinstance(python_config_type, type) and issubclass( + python_config_type, Enum + ): return python_config_type[enum_name] return enum_name @@ -170,6 +178,9 @@ def main() -> None: java_agent_execution_options = class_loader.loadClass( "org.apache.flink.agents.api.agents.AgentExecutionOptions" ) + java_memory_event_options = class_loader.loadClass( + "org.apache.flink.agents.api.configuration.MemoryEventOptions" + ) assert_options_class_matches_java( AgentConfigOptions, java_agent_config_options, jvm @@ -177,6 +188,9 @@ def main() -> None: assert_options_class_matches_java( AgentExecutionOptions, java_agent_execution_options, jvm ) + assert_options_class_matches_java( + MemoryEventOptions, java_memory_event_options, jvm + ) if __name__ == "__main__": diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index ada5b8e309..2384dd6ea3 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -805,19 +805,6 @@ def _init_long_term_memory( ) -def flink_runner_context_switch_action_context( - ctx: FlinkRunnerContext, - key: int, -) -> None: - """Switch the context of the flink runner context. - - The ctx is reused across keyed partitions, the context related to - specific key should be switched when process new action. - """ - if ctx.long_term_memory is not None: - ctx.long_term_memory.switch_context(str(key)) - - def close_flink_runner_context( ctx: FlinkRunnerContext, ) -> None: diff --git a/python/flink_agents/runtime/memory/internal_base_long_term_memory.py b/python/flink_agents/runtime/memory/internal_base_long_term_memory.py index 96ff4558ac..3050f4ca30 100644 --- a/python/flink_agents/runtime/memory/internal_base_long_term_memory.py +++ b/python/flink_agents/runtime/memory/internal_base_long_term_memory.py @@ -26,11 +26,29 @@ class InternalBaseLongTermMemory(BaseLongTermMemory, ABC): """ @abstractmethod - def switch_context(self, key: str) -> None: + def configure_observation( + self, + *, + update_observation_enabled: bool, + get_observation_enabled: bool, + search_observation_enabled: bool, + ) -> None: + """Configure which operations produce observation records.""" + + @abstractmethod + def switch_context( + self, + key: str, + *, + observation_id: str, + observation_suppressed: bool = False, + ) -> None: """Switches the context for the memory operations. This allows the same memory instance to be used for different key by isolating data based on the provided key. Args: key: The context key. + observation_id: Identifier for the current action's observations. + observation_suppressed: Whether to suppress observation for the action. """ 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 bcd81dfd32..db87a3470d 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 @@ -18,9 +18,12 @@ from __future__ import annotations import contextlib +import json import logging import queue +from dataclasses import dataclass from datetime import datetime +from enum import Enum from typing import Any, Dict, List from pydantic import ConfigDict, Field, PrivateAttr, field_validator @@ -46,6 +49,42 @@ _PROVIDER_NAME = "flink_agents" +class _LtmObservationOp(str, Enum): + ADD = "ADD" + UPDATE = "UPDATE" + DELETE = "DELETE" + DELETE_SET = "DELETE_SET" + GET = "GET" + SEARCH = "SEARCH" + + +_MEM0_ADD_RESULT_OPERATIONS = { + "ADD": _LtmObservationOp.ADD, + "UPDATE": _LtmObservationOp.UPDATE, + "DELETE": _LtmObservationOp.DELETE, +} + + +@dataclass(frozen=True) +class _LtmObservationRecord: + op: str + set: str + id: str | None = None + query: str | None = None + value: Any = None + version: int = 1 + + def to_wire(self) -> Dict[str, Any]: + return { + "op": self.op, + "set": self.set, + "id": self.id, + "query": self.query, + "value": self.value, + "version": self.version, + } + + def _create_flink_agents_config_classes() -> tuple: """Create custom LLM/Embedder config classes that accept the ``flink_agents`` provider. @@ -128,6 +167,15 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory): _mem0: Any = PrivateAttr(default=None) + _ltm_observation_records: queue.Queue[tuple[str, str, _LtmObservationRecord]] = ( + PrivateAttr(default_factory=queue.Queue) + ) + _update_observation_enabled: bool = PrivateAttr(default=False) + _get_observation_enabled: bool = PrivateAttr(default=False) + _search_observation_enabled: bool = PrivateAttr(default=False) + _observation_id: str = PrivateAttr(default="") + _observation_suppressed: bool = PrivateAttr(default=False) + def __init__( self, *, @@ -224,7 +272,26 @@ def _create_mem0_instance(self) -> Any: return Memory(mem0_config) @override - def switch_context(self, key: str) -> None: + def configure_observation( + self, + *, + update_observation_enabled: bool, + get_observation_enabled: bool, + search_observation_enabled: bool, + ) -> None: + """Configure the fixed per-operation observation switches.""" + self._update_observation_enabled = update_observation_enabled + self._get_observation_enabled = get_observation_enabled + self._search_observation_enabled = search_observation_enabled + + @override + def switch_context( + self, + key: str, + *, + observation_id: str, + observation_suppressed: bool = False, + ) -> None: """Switch the keyed partition context. This method is called on the mailbox thread before each action. @@ -233,12 +300,130 @@ def switch_context(self, key: str) -> None: Args: key: The new key for partition isolation. + observation_id: Identifier for the current action's observations. + 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_id = observation_id + 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, + observation_id: 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. + observation_id: Action identifier 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, + observation_id, + _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, + observation_id: 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. + observation_id: Action identifier captured at operation entry. + enabled: Whether search observation is configured. + """ + try: + if not enabled: + return + self._ltm_observation_records.put( + ( + observation_key, + observation_id, + _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, observation_id: str) -> str: + """Pop buffered LTM records for one action as a JSON array. + + Called from Java on the mailbox thread at action-finish flush. Records owned + by other partition/action pairs are placed back in the shared queue. + + Args: + key: Partition key whose records to drain. + observation_id: Action identifier whose records to drain. + + Returns: + JSON array string of the drained records. + """ + records: List[_LtmObservationRecord] = [] + other_records: List[tuple[str, str, _LtmObservationRecord]] = [] + while True: + try: + owner_key, owner_observation_id, record = ( + self._ltm_observation_records.get_nowait() + ) + except queue.Empty: + break + if owner_key == key and owner_observation_id == observation_id: + records.append(record) + else: + other_records.append((owner_key, owner_observation_id, record)) + for record in other_records: + self._ltm_observation_records.put(record) + + return json.dumps([record.to_wire() for record in records], ensure_ascii=False) @override def get_memory_set(self, name: str) -> MemorySet: @@ -262,11 +447,25 @@ def delete_memory_set(self, name: str) -> bool: Returns: True if the memory set was deleted. """ + observation_key = self.key + observation_id = self._observation_id + observation_enabled = ( + self._update_observation_enabled and not self._observation_suppressed + ) self._mem0_instance.delete_all( user_id=self.job_id, - agent_id=self.key, + agent_id=observation_key, run_id=name, ) + self._record_ltm_op( + _LtmObservationOp.DELETE_SET, + name, + None, + None, + observation_key, + observation_id, + enabled=observation_enabled, + ) return True @override @@ -286,6 +485,11 @@ def add( Returns: List of IDs of the added memories. """ + observation_key = self.key + observation_id = self._observation_id + observation_enabled = ( + self._update_observation_enabled and not self._observation_suppressed + ) if isinstance(memory_items, str): memory_items = [memory_items] if metadatas is not None and isinstance(metadatas, dict): @@ -297,14 +501,33 @@ def add( result = self._mem0_instance.add( messages=item, user_id=self.job_id, - agent_id=self.key if self.key else None, + agent_id=observation_key, run_id=memory_set.name, metadata=metadata, ) # Extract IDs from the result - all_ids.extend( - entry["id"] for entry in result.get("results", []) if "id" in entry - ) + for entry in result.get("results", []): + if "id" in entry: + all_ids.append(entry["id"]) + event_name = str(entry.get("event", "")).upper() + operation = _MEM0_ADD_RESULT_OPERATIONS.get(event_name) + if operation is None: + logger.warning( + "Skipping unsupported Mem0 add result operation %r", + event_name, + ) + else: + self._record_ltm_op( + operation, + memory_set.name, + entry["id"], + None + if operation == _LtmObservationOp.DELETE + else entry.get("memory"), + observation_key, + observation_id, + enabled=observation_enabled, + ) return all_ids @override @@ -327,6 +550,11 @@ def get( Returns: List of memory items. """ + observation_key = self.key + observation_id = self._observation_id + observation_enabled = ( + self._get_observation_enabled and not self._observation_suppressed + ) if ids is not None: if isinstance(ids, str): ids = [ids] @@ -334,19 +562,40 @@ def get( for memory_id in ids: result = self._mem0_instance.get(memory_id=memory_id) items.append(self._convert_mem0_result(memory_set.name, result)) + item = items[-1] + self._record_ltm_op( + _LtmObservationOp.GET, + memory_set.name, + item.id, + item.value, + observation_key, + observation_id, + enabled=observation_enabled, + ) return items result = self._mem0_instance.get_all( user_id=self.job_id, - agent_id=self.key, + agent_id=observation_key, run_id=memory_set.name, filters=filters, limit=limit, ) - return [ + items = [ self._convert_mem0_result(memory_set.name, entry) for entry in result.get("results", []) ] + for item in items: + self._record_ltm_op( + _LtmObservationOp.GET, + memory_set.name, + item.id, + item.value, + observation_key, + observation_id, + enabled=observation_enabled, + ) + return items @override def delete(self, memory_set: MemorySet, ids: str | List[str] | None = None) -> None: @@ -356,18 +605,41 @@ 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_enabled = ( + self._update_observation_enabled and not self._observation_suppressed + ) if ids is None: self._mem0_instance.delete_all( user_id=self.job_id, - agent_id=self.key, + agent_id=observation_key, run_id=memory_set.name, ) + self._record_ltm_op( + _LtmObservationOp.DELETE_SET, + memory_set.name, + None, + None, + observation_key, + observation_id, + enabled=observation_enabled, + ) return if isinstance(ids, str): ids = [ids] for memory_id in ids: self._mem0_instance.delete(memory_id=memory_id) + self._record_ltm_op( + _LtmObservationOp.DELETE, + memory_set.name, + memory_id, + None, + observation_key, + observation_id, + enabled=observation_enabled, + ) @override def search( @@ -390,18 +662,35 @@ def search( Returns: List of matching memory items. """ + observation_key = self.key + observation_id = self._observation_id + observation_enabled = ( + self._search_observation_enabled and not self._observation_suppressed + ) result = self._mem0_instance.search( query=query, user_id=self.job_id, - agent_id=self.key, + agent_id=observation_key, run_id=memory_set.name, limit=limit, filters=filters, **kwargs, ) + raw_entries = result.get("results", []) + self._record_ltm_search( + memory_set.name, + query, + [ + {"id": e.get("id"), "value": e.get("memory"), "score": e.get("score")} + for e in raw_entries + if "id" in e + ], + observation_key, + observation_id, + enabled=observation_enabled, + ) return [ - self._convert_mem0_result(memory_set.name, entry) - for entry in result.get("results", []) + self._convert_mem0_result(memory_set.name, entry) for entry in raw_entries ] @staticmethod @@ -455,7 +744,9 @@ def _report_token_metrics(self) -> None: and metric.get("promptTokens") and metric.get("completionTokens") ): - model_group = self.metric_group.get_sub_group("model", metric["model_name"]) + model_group = self.metric_group.get_sub_group( + "model", metric["model_name"] + ) model_group.get_counter("promptTokens").inc(metric["promptTokens"]) model_group.get_counter("completionTokens").inc( metric["completionTokens"] 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 9911bb4464..3f3faf4e59 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 @@ -325,11 +325,11 @@ def test_delete_memory_set(ltm) -> None: def test_switch_context(ltm) -> None: - ltm.switch_context("key_a") + ltm.switch_context("key_a", observation_id="action-a") memory_set = ltm.get_memory_set(name="context_set") memory_set.add(items="Data for key_a") - ltm.switch_context("key_b") + 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 @@ -337,7 +337,7 @@ def test_switch_context(ltm) -> None: assert len(items) == 0 # Reset context - ltm.switch_context("") + ltm.switch_context("", observation_id="action-empty") class MockChatModelWithTokenUsage: @@ -462,7 +462,7 @@ def get_resource(name: str, resource_type: Any, **kwargs: Any) -> Any: ) # First switch_context triggers lazy init; no metrics yet. - ltm.switch_context("key_1") + ltm.switch_context("key_1", observation_id="action-1") memory_set = ltm.get_memory_set(name="token_test_set") memory_set.add(items="Test content for token tracking.") @@ -473,7 +473,7 @@ def get_resource(name: str, resource_type: Any, **kwargs: Any) -> Any: assert model_group._counters.get("promptTokens", 0) == 0 # switch_context drains the queue on the mailbox thread. - ltm.switch_context("key_2") + ltm.switch_context("key_2", observation_id="action-2") # Mem0 makes 2 LLM calls per add (fact extraction + memory update). # Each call returns 10 prompt tokens and 5 completion tokens. @@ -521,7 +521,7 @@ def get_resource(name: str, resource_type: Any, **kwargs: Any) -> Any: vector_store_name="test_vector_store", ) - ltm.switch_context("key_1") + ltm.switch_context("key_1", observation_id="action-1") memory_set = ltm.get_memory_set(name="close_test_set") memory_set.add(items="First item.") 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 new file mode 100644 index 0000000000..4b83aec254 --- /dev/null +++ b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py @@ -0,0 +1,214 @@ +################################################################################ +# 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. +################################################################################# +"""Public Mem0 operations produce correctly typed observations.""" + +import json +from concurrent.futures import ThreadPoolExecutor +from inspect import signature +from threading import Event +from typing import Any +from unittest.mock import MagicMock + +from flink_agents.runtime.memory.internal_base_long_term_memory import ( + InternalBaseLongTermMemory, +) +from flink_agents.runtime.memory.mem0.mem0_long_term_memory import Mem0LongTermMemory + + +def _make_ltm(mem0: Any) -> Mem0LongTermMemory: + ctx = MagicMock() + ctx.agent_metric_group = None + ltm = Mem0LongTermMemory.model_construct( + ctx=ctx, job_id="job", key="partition", metric_group=None + ) + ltm._mem0 = mem0 + ltm._observation_id = "action" + ltm.configure_observation( + update_observation_enabled=True, + get_observation_enabled=True, + search_observation_enabled=True, + ) + return ltm + + +def _drain( + ltm: Mem0LongTermMemory, + key: str = "partition", + observation_id: str = "action", +) -> list[dict]: + return json.loads(ltm.drain_ltm_observation_records(key, observation_id)) + + +def test_internal_context_signatures_match_implementation() -> None: + assert signature( + InternalBaseLongTermMemory.configure_observation, eval_str=True + ) == signature( + Mem0LongTermMemory.configure_observation, + eval_str=True, + ) + assert signature( + InternalBaseLongTermMemory.switch_context, eval_str=True + ) == signature( + Mem0LongTermMemory.switch_context, + eval_str=True, + ) + + +def test_add_maps_mem0_add_update_delete_results() -> None: + mem0 = MagicMock() + mem0.add.return_value = { + "results": [ + {"event": "ADD", "id": "a", "memory": "added"}, + {"event": "UPDATE", "id": "u", "memory": "updated"}, + {"event": "DELETE", "id": "d", "memory": "ignored"}, + ] + } + ltm = _make_ltm(mem0) + + assert ltm.add(ltm.get_memory_set("prefs"), "input") == ["a", "u", "d"] + assert [ + (record["op"], record["id"], record["value"]) for record in _drain(ltm) + ] == [ + ("ADD", "a", "added"), + ("UPDATE", "u", "updated"), + ("DELETE", "d", None), + ] + + +def test_unknown_mem0_result_is_not_mislabeled_as_add() -> None: + mem0 = MagicMock() + mem0.add.return_value = { + "results": [{"event": "NOOP", "id": "x", "memory": "value"}] + } + ltm = _make_ltm(mem0) + assert ltm.add(ltm.get_memory_set("prefs"), "input") == ["x"] + assert _drain(ltm) == [] + + +def test_public_method_captures_observation_owner_at_entry() -> None: + mem0 = MagicMock() + ltm = _make_ltm(mem0) + memory_set = ltm.get_memory_set("prefs") + + def finish_after_context_switch(**_kwargs: Any) -> dict: + ltm.key = "partition-2" + ltm._observation_id = "action-2" + return {"results": [{"event": "ADD", "id": "m1", "memory": "v"}]} + + mem0.add.side_effect = finish_after_context_switch + ltm.add(memory_set, "input") + + assert [record["id"] for record in _drain(ltm)] == ["m1"] + assert _drain(ltm, "partition-2", "action-2") == [] + + +def test_add_uses_entry_context_when_another_action_switches_context() -> None: + mem0 = MagicMock() + mem0.add.return_value = {"results": []} + ltm = _make_ltm(mem0) + memory_set = ltm.get_memory_set("prefs") + observation_config_read = Event() + resume_operation = Event() + + class BlockingEnabledFlag: + def __bool__(self) -> bool: + observation_config_read.set() + if not resume_operation.wait(timeout=5): + message = "Timed out waiting for the context switch" + raise TimeoutError(message) + return True + + ltm._update_observation_enabled = BlockingEnabledFlag() + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(ltm.add, memory_set, "input") + assert observation_config_read.wait(timeout=5) + ltm.switch_context("partition-2", observation_id="action-2") + resume_operation.set() + future.result(timeout=5) + + assert mem0.add.call_args.kwargs["agent_id"] == "partition" + + +def test_get_delete_and_search_keep_structured_identity() -> None: + mem0 = MagicMock() + mem0.get.return_value = {"id": "g1", "memory": "value"} + mem0.search.return_value = { + "results": [{"id": "s1", "memory": "hit", "score": 0.9}] + } + ltm = _make_ltm(mem0) + memory_set = ltm.get_memory_set("policies") + + ltm.get(memory_set, ids="g1") + ltm.delete(memory_set, ids="d1") + ltm.search(memory_set, "refund", limit=5) + + records = _drain(ltm) + assert [(record["op"], record["set"]) for record in records] == [ + ("GET", "policies"), + ("DELETE", "policies"), + ("SEARCH", "policies"), + ] + assert records[2]["query"] == "refund" + + +def test_context_switch_changes_observation_owner_and_current_suppression() -> None: + mem0 = MagicMock() + mem0.add.return_value = { + "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") + assert _drain(ltm, "suppressed", "suppressed-action") == [] + + ltm.switch_context("observed", observation_id="observed-action") + ltm.add(memory_set, "recorded") + assert [record["id"] for record in _drain(ltm, "observed", "observed-action")] == [ + "m1" + ] + assert ltm._update_observation_enabled is True + assert ltm._get_observation_enabled is True + assert ltm._search_observation_enabled is True + + +def test_empty_context_key_is_used_consistently_for_mem0_operations() -> None: + mem0 = MagicMock() + mem0.add.return_value = {"results": []} + 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") + ltm.add(memory_set, "input") + ltm.get(memory_set) + ltm.search(memory_set, "query", limit=5) + ltm.delete(memory_set) + ltm.delete_memory_set("prefs") + + assert mem0.add.call_args.kwargs["agent_id"] == "" + assert mem0.get_all.call_args.kwargs["agent_id"] == "" + assert mem0.search.call_args.kwargs["agent_id"] == "" + assert [call.kwargs["agent_id"] for call in mem0.delete_all.call_args_list] == [ + "", + "", + ] diff --git a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_recording_hook.py b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_recording_hook.py new file mode 100644 index 0000000000..765efe7047 --- /dev/null +++ b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_recording_hook.py @@ -0,0 +1,172 @@ +################################################################################ +# 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. +################################################################################# +"""Tests for the private typed LTM observation buffer.""" + +import json +import threading +from typing import Any +from unittest.mock import MagicMock + +from flink_agents.runtime.memory.mem0.mem0_long_term_memory import ( + Mem0LongTermMemory, + _LtmObservationOp, +) + + +def _make_ltm() -> Mem0LongTermMemory: + ctx = MagicMock() + ctx.agent_metric_group = None + return Mem0LongTermMemory.model_construct( + ctx=ctx, job_id="job", key="shared-partition", metric_group=None + ) + + +def test_records_are_versioned_and_search_keeps_set_and_query() -> None: + ltm = _make_ltm() + ltm._record_ltm_op(_LtmObservationOp.ADD, "prefs", "m1", "v", "partition", "action") + ltm._record_ltm_search("policies", "refund", [{"id": "p1"}], "partition", "action") + + assert json.loads(ltm.drain_ltm_observation_records("partition", "action")) == [ + { + "op": "ADD", + "set": "prefs", + "id": "m1", + "query": None, + "value": "v", + "version": 1, + }, + { + "op": "SEARCH", + "set": "policies", + "id": None, + "query": "refund", + "value": [{"id": "p1"}], + "version": 1, + }, + ] + + +def test_drain_matches_partition_and_observation_id() -> None: + ltm = _make_ltm() + ltm._record_ltm_op( + _LtmObservationOp.ADD, + "s", + "partition-a-record", + "va", + "partition-a", + "shared-action", + ) + ltm._record_ltm_op( + _LtmObservationOp.ADD, + "s", + "partition-b-record", + "vb", + "partition-b", + "shared-action", + ) + ltm._record_ltm_op( + _LtmObservationOp.ADD, + "s", + "action-a-record", + "va", + "shared-partition", + "action-a", + ) + ltm._record_ltm_op( + _LtmObservationOp.ADD, + "s", + "action-b-record", + "vb", + "shared-partition", + "action-b", + ) + + assert [ + record["id"] + for record in json.loads( + ltm.drain_ltm_observation_records("partition-b", "shared-action") + ) + ] == ["partition-b-record"] + assert [ + record["id"] + for record in json.loads( + ltm.drain_ltm_observation_records("shared-partition", "action-b") + ) + ] == ["action-b-record"] + assert [ + record["id"] + for record in json.loads( + ltm.drain_ltm_observation_records("partition-a", "shared-action") + ) + ] == ["partition-a-record"] + assert [ + record["id"] + for record in json.loads( + ltm.drain_ltm_observation_records("shared-partition", "action-a") + ) + ] == ["action-a-record"] + + +def test_disabled_operation_does_not_enqueue_or_json_encode() -> None: + ltm = _make_ltm() + ltm._record_ltm_op( + _LtmObservationOp.ADD, + "s", + "m1", + object(), + "partition", + "action", + enabled=False, + ) + assert ltm._ltm_observation_records.empty() + + +def test_buffer_is_thread_safe() -> None: + ltm = _make_ltm() + + def worker() -> None: + for index in range(100): + ltm._record_ltm_op( + _LtmObservationOp.ADD, + "s", + f"id-{index}", + "v", + "partition", + "action", + ) + + threads = [threading.Thread(target=worker) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert ( + len(json.loads(ltm.drain_ltm_observation_records("partition", "action"))) == 200 + ) + + +class _RaisingQueue: + def put(self, _item: Any) -> None: + message = "boom" + raise RuntimeError(message) + + +def test_buffering_hook_never_raises() -> None: + ltm = _make_ltm() + ltm._ltm_observation_records = _RaisingQueue() + ltm._record_ltm_op(_LtmObservationOp.ADD, "s", "m1", "v", "partition", "action") diff --git a/python/flink_agents/runtime/python_java_utils.py b/python/flink_agents/runtime/python_java_utils.py index 3c4bb60dcf..fc1b066622 100644 --- a/python/flink_agents/runtime/python_java_utils.py +++ b/python/flink_agents/runtime/python_java_utils.py @@ -49,6 +49,13 @@ def convert_to_python_object(bytesObject: bytes) -> Any: return cloudpickle.loads(bytesObject) +def convert_to_python_key_text(bytes_object: bytes, serialization: str) -> str: + """Return the textual identity for a serialized or explicitly typed bytes key.""" + if serialization == "pickled": + return str(cloudpickle.loads(bytes_object)) + return str(bytes_object) + + def convert_json_to_python_event(event_json: str) -> Event: """Deserialize a JSON string into a base Python Event object. 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 3d6b0decba..20dce808b4 100644 --- a/python/flink_agents/runtime/tests/test_python_java_utils.py +++ b/python/flink_agents/runtime/tests/test_python_java_utils.py @@ -28,6 +28,7 @@ from flink_agents.api.tools import InjectedArg from flink_agents.runtime.python_java_utils import ( call_embedding_with_usage, + convert_to_python_key_text, get_python_tool_metadata, wrap_to_input_event, ) @@ -82,3 +83,19 @@ def test_wrap_to_input_event_assigns_unique_random_id_per_record() -> None: assert first.id != second.id assert first.id.version == 4 assert second.id.version == 4 + + +def test_convert_to_python_key_text_uses_python_str() -> None: + assert convert_to_python_key_text(cloudpickle.dumps(7), "pickled") == "7" + assert ( + convert_to_python_key_text( + cloudpickle.dumps(("tenant", 7)), + "pickled", + ) + == "('tenant', 7)" + ) + + +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.'" diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java b/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java index 753fd7ea8d..37351f2cf9 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/CompileUtils.java @@ -24,10 +24,14 @@ import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory; import org.apache.flink.api.common.typeinfo.TypeInformation; import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.typeutils.RowTypeInfo; import org.apache.flink.streaming.api.datastream.DataStream; import org.apache.flink.streaming.api.datastream.KeyedStream; +import org.apache.flink.streaming.api.typeinfo.python.PickledByteArrayTypeInfo; import org.apache.flink.types.Row; +import static org.apache.flink.util.Preconditions.checkArgument; + /** A utility class that bridges Flink DataStream/SQL with the Flink Agents agent. */ public class CompileUtils { @@ -37,7 +41,12 @@ public static DataStream connectToAgent( throws JsonProcessingException { // deserialize agent plan json. AgentPlan agentPlan = new ObjectMapper().readValue(agentPlanJson, AgentPlan.class); - return connectToAgent(inputDataStream, agentPlan, TypeInformation.of(byte[].class), false); + return connectToAgent( + inputDataStream, + agentPlan, + TypeInformation.of(byte[].class), + false, + isPickledPythonKeyType(inputDataStream.getKeyType())); } // ============================ invoke by java ==================================== @@ -48,7 +57,8 @@ public static DataStream connectToAgent( public static DataStream connectToAgent( KeyedStream keyedInputStream, AgentPlan agentPlan) { - return connectToAgent(keyedInputStream, agentPlan, TypeInformation.of(Object.class), true); + return connectToAgent( + keyedInputStream, agentPlan, TypeInformation.of(Object.class), true, false); } // ============================ basic ==================================== @@ -73,13 +83,32 @@ private static DataStream connectToAgent( KeyedStream keyedInputStream, AgentPlan agentPlan, TypeInformation outTypeInformation, - boolean inputIsJava) { + boolean inputIsJava, + boolean pythonKeyIsPickled) { return (DataStream) keyedInputStream .transform( "action-execute-operator", outTypeInformation, - new ActionExecutionOperatorFactory(agentPlan, inputIsJava)) + new ActionExecutionOperatorFactory( + agentPlan, inputIsJava, pythonKeyIsPickled)) .setParallelism(keyedInputStream.getParallelism()); } + + /** + * Returns whether PyFlink's single logical key field uses its default pickle representation. + */ + static boolean isPickledPythonKeyType(TypeInformation keyType) { + checkArgument( + keyType instanceof RowTypeInfo, + "Expected PyFlink key type to be a single-field RowTypeInfo, but got %s", + keyType); + RowTypeInfo rowType = (RowTypeInfo) keyType; + checkArgument( + rowType.getArity() == 1, + "Expected PyFlink key type to contain one logical field, but got arity %s", + rowType.getArity()); + TypeInformation logicalKeyType = rowType.getTypeAt(0); + return logicalKeyType instanceof PickledByteArrayTypeInfo; + } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java index 1395bdc568..2049e1e021 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java @@ -38,7 +38,10 @@ import org.apache.flink.agents.runtime.actionstate.CallResult; import org.apache.flink.agents.runtime.memory.CachedMemoryStore; import org.apache.flink.agents.runtime.memory.InteranlBaseLongTermMemory; +import org.apache.flink.agents.runtime.memory.MemoryEventBuilder; +import org.apache.flink.agents.runtime.memory.MemoryEventSettings; import org.apache.flink.agents.runtime.memory.MemoryObjectImpl; +import org.apache.flink.agents.runtime.memory.MemoryValueObservation; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import org.apache.flink.util.Preconditions; import org.slf4j.Logger; @@ -47,6 +50,7 @@ import javax.annotation.Nullable; import java.util.ArrayList; +import java.util.Collections; import java.util.LinkedList; import java.util.List; import java.util.Map; @@ -67,6 +71,8 @@ public static class MemoryContext { private final CachedMemoryStore shortTermMemStore; private final List sensoryMemoryUpdates; private final List shortTermMemoryUpdates; + private final List sensoryMemoryReads; + private final List shortTermMemoryReads; public MemoryContext( CachedMemoryStore sensoryMemStore, CachedMemoryStore shortTermMemStore) { @@ -74,6 +80,8 @@ public MemoryContext( this.shortTermMemStore = shortTermMemStore; this.sensoryMemoryUpdates = new LinkedList<>(); this.shortTermMemoryUpdates = new LinkedList<>(); + this.sensoryMemoryReads = new LinkedList<>(); + this.shortTermMemoryReads = new LinkedList<>(); } public List getShortTermMemoryUpdates() { @@ -84,6 +92,19 @@ public List getSensoryMemoryUpdates() { return sensoryMemoryUpdates; } + public List getSensoryMemoryReads() { + return sensoryMemoryReads; + } + + public List getShortTermMemoryReads() { + return shortTermMemoryReads; + } + + private void clearReadObservations() { + sensoryMemoryReads.clear(); + shortTermMemoryReads.clear(); + } + public CachedMemoryStore getShortTermMemStore() { return shortTermMemStore; } @@ -105,6 +126,24 @@ public CachedMemoryStore getSensoryMemStore() { protected String actionName; protected InteranlBaseLongTermMemory ltm; + /** Textual key shared by long-term-memory isolation and framework observation events. */ + protected String contextKey; + + /** Stable identifier that isolates observations for one logical action execution. */ + protected String observationId; + + /** True when the current action was triggered by a memory event: suppress observation. */ + protected boolean observationSuppressed; + + /** True when at least one LTM operation records observations for the current action. */ + protected boolean ltmObservationEnabled; + + /** Resolved per-operation memory-event switches; config is fixed per agent plan. */ + private final MemoryEventSettings memoryEventSettings; + + /** Whether the fixed job-level configuration enables any LTM observation. */ + private final boolean ltmObservationConfigured; + /** Context for fine-grained durable execution, may be null if not enabled. */ @Nullable protected DurableExecutionContext durableExecutionContext; @@ -118,17 +157,32 @@ public RunnerContextImpl( this.mailboxThreadChecker = mailboxThreadChecker; this.agentPlan = agentPlan; this.resourceCache = resourceCache; + this.memoryEventSettings = MemoryEventSettings.from(agentPlan.getConfigData()); + this.ltmObservationConfigured = + memoryEventSettings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_UPDATE) + || memoryEventSettings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_GET) + || memoryEventSettings.generate( + MemoryEventSettings.MemoryOp.LONG_TERM_SEARCH); } 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 contextKey, + String observationId, + boolean observationSuppressed) { this.actionName = actionName; this.memoryContext = memoryContext; + this.contextKey = contextKey; + this.observationId = observationId; + this.observationSuppressed = observationSuppressed; + this.ltmObservationEnabled = !observationSuppressed && ltmObservationConfigured; if (ltm != null) { - ltm.switchContext(key); + ltm.switchContext(contextKey, observationId, observationSuppressed); } } @@ -161,6 +215,32 @@ public void sendEvent(Event event) { public List drainEvents(Long timestamp) { mailboxThreadChecker.run(); + return drainPendingEvents(timestamp); + } + + /** Converts this action's memory records into events and drains all action output events. */ + public List drainEventsAtActionFinish(Long timestamp) { + mailboxThreadChecker.run(); + flushMemoryObservation(); + return drainPendingEvents(timestamp); + } + + /** + * Discards pending LTM observation records for the current key without rolling back written + * data. + */ + public void discardMemoryObservation() { + mailboxThreadChecker.run(); + if (memoryContext != null) { + memoryContext.clearReadObservations(); + } + if (ltm == null || !ltmObservationEnabled) { + return; + } + ltm.drainObservationRecordsJson(contextKey, observationId); + } + + private List drainPendingEvents(Long timestamp) { List list = new ArrayList<>(this.pendingEvents); if (timestamp != null) { list.forEach(event -> event.setSourceTimestamp(timestamp)); @@ -169,6 +249,52 @@ public List drainEvents(Long timestamp) { return list; } + private void flushMemoryObservation() { + if (memoryContext == null) { + return; + } + List sensoryReads = + new ArrayList<>(memoryContext.getSensoryMemoryReads()); + List shortTermReads = + new ArrayList<>(memoryContext.getShortTermMemoryReads()); + memoryContext.clearReadObservations(); + if (observationSuppressed || !memoryEventSettings.anyEnabled()) { + return; + } + List> ltmRecords = Collections.emptyList(); + if (ltm != null && ltmObservationEnabled) { + try { + ltmRecords = + MemoryEventBuilder.parseLtmObservationRecords( + ltm.drainObservationRecordsJson(contextKey, observationId)); + } catch (Exception | LinkageError e) { + LOG.warn( + "LTM observation drain failed for action '{}' and partition key '{}' ({}); skipping records", + actionName, + contextKey, + e.getClass().getSimpleName()); + } + } + try { + pendingEvents.addAll( + MemoryEventBuilder.buildWriteEvents( + contextKey, + memoryContext.getSensoryMemoryUpdates(), + memoryContext.getShortTermMemoryUpdates(), + memoryEventSettings)); + pendingEvents.addAll( + MemoryEventBuilder.buildReadEvents( + contextKey, sensoryReads, shortTermReads, memoryEventSettings)); + pendingEvents.addAll( + MemoryEventBuilder.buildLtmEvents(contextKey, ltmRecords, memoryEventSettings)); + } catch (RuntimeException | LinkageError e) { + LOG.warn( + "Skipping framework memory observation for action '{}' ({})", + actionName, + e.getClass().getSimpleName()); + } + } + public void checkNoPendingEvents() { Preconditions.checkState( this.pendingEvents.isEmpty(), "There are pending events remaining in the context."); @@ -194,23 +320,35 @@ public List getShortTermMemoryUpdates() { @Override public MemoryObject getSensoryMemory() throws Exception { mailboxThreadChecker.run(); + List memoryReads = null; + if (!observationSuppressed + && memoryEventSettings.generate(MemoryEventSettings.MemoryOp.SENSORY_READ)) { + memoryReads = memoryContext.getSensoryMemoryReads(); + } return new MemoryObjectImpl( MemoryObject.MemoryType.SENSORY, memoryContext.getSensoryMemStore(), MemoryObjectImpl.ROOT_KEY, mailboxThreadChecker, - memoryContext.getSensoryMemoryUpdates()); + memoryContext.getSensoryMemoryUpdates(), + memoryReads); } @Override public MemoryObject getShortTermMemory() throws Exception { mailboxThreadChecker.run(); + List memoryReads = null; + if (!observationSuppressed + && memoryEventSettings.generate(MemoryEventSettings.MemoryOp.SHORT_TERM_READ)) { + memoryReads = memoryContext.getShortTermMemoryReads(); + } return new MemoryObjectImpl( MemoryObject.MemoryType.SHORT_TERM, memoryContext.getShortTermMemStore(), MemoryObjectImpl.ROOT_KEY, mailboxThreadChecker, - memoryContext.getShortTermMemoryUpdates()); + memoryContext.getShortTermMemoryUpdates(), + memoryReads); } @Override diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/InteranlBaseLongTermMemory.java b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/InteranlBaseLongTermMemory.java index b13e2d7ca4..25748be233 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/InteranlBaseLongTermMemory.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/InteranlBaseLongTermMemory.java @@ -21,11 +21,34 @@ /** Internal interface extends {@link BaseLongTermMemory} for hiding some interface to user. */ public interface InteranlBaseLongTermMemory extends BaseLongTermMemory { + /** + * Configures which long-term-memory operations produce observation records. + * + *

The configuration is fixed for an agent plan and should be applied once when the backend + * is wired. + */ + void configureObservation( + boolean updateObservationEnabled, + boolean getObservationEnabled, + boolean searchObservationEnabled); + /** * Switches the context for the memory operations. This allows the same memory instance to be * used for different key by isolating data based on the provided key. * - * @param key the context key + * @param partitionKey the context key used by the long-term-memory backend + * @param observationId the action-scoped identifier used only for observation isolation + * @param observationSuppressed whether observation records are suppressed for this action + */ + void switchContext(String partitionKey, String observationId, boolean observationSuppressed); + + /** + * Drains the LTM observation records buffered for one action scope, returning them as a JSON + * array string. Called on the mailbox thread at the action finish boundary. + * + * @param partitionKey the partition whose records to drain + * @param observationId the action whose records to drain + * @return JSON array string of drained records */ - void switchContext(String key); + String drainObservationRecordsJson(String partitionKey, String observationId); } 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 2847812a44..11d95ac8dd 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 @@ -130,8 +130,41 @@ public List search( } @Override - public void switchContext(String key) { - adapter.callMethod(pyMem0, "switch_context", Map.of("key", key)); + public void configureObservation( + boolean updateObservationEnabled, + boolean getObservationEnabled, + boolean searchObservationEnabled) { + adapter.callMethod( + pyMem0, + "configure_observation", + Map.of( + "update_observation_enabled", + updateObservationEnabled, + "get_observation_enabled", + getObservationEnabled, + "search_observation_enabled", + searchObservationEnabled)); + } + + @Override + public void switchContext( + String partitionKey, String observationId, boolean observationSuppressed) { + adapter.callMethod( + pyMem0, + "switch_context", + Map.of( + "key", partitionKey, + "observation_id", observationId, + "observation_suppressed", observationSuppressed)); + } + + @Override + public String drainObservationRecordsJson(String partitionKey, String observationId) { + return (String) + adapter.callMethod( + pyMem0, + "drain_ltm_observation_records", + Map.of("key", partitionKey, "observation_id", observationId)); } @Override diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryEventBuilder.java b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryEventBuilder.java new file mode 100644 index 0000000000..a1839c680b --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryEventBuilder.java @@ -0,0 +1,341 @@ +/* + * 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.runtime.memory; + +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.context.MemoryUpdate; +import org.apache.flink.agents.api.event.LongTermUpdateEvent; +import org.apache.flink.agents.api.event.MemoryEvent; +import org.apache.flink.agents.runtime.memory.MemoryEventSettings.MemoryOp; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Function; +import java.util.function.Supplier; + +/** Converts per-action memory observations into at most one event of each operation type. */ +public final class MemoryEventBuilder { + private static final Logger LOG = LoggerFactory.getLogger(MemoryEventBuilder.class); + private static final ObjectMapper MAPPER = + JsonMapper.builder().disable(JsonWriteFeature.WRITE_NAN_AS_STRINGS).build(); + private static final int LTM_RECORD_VERSION = 1; + private static final Object INVALID_VALUE = new Object(); + + private enum LtmOperation { + ADD, + UPDATE, + DELETE, + DELETE_SET, + GET, + SEARCH + } + + private MemoryEventBuilder() {} + + public static List buildWriteEvents( + String eventKeyText, + List sensoryWrites, + List shortTermWrites, + MemoryEventSettings settings) { + List events = new ArrayList<>(2); + addFoldedEvent( + events, + eventKeyText, + MemoryOp.SENSORY_WRITE, + sensoryWrites, + MemoryUpdate::getPath, + MemoryUpdate::getValue, + settings); + addFoldedEvent( + events, + eventKeyText, + MemoryOp.SHORT_TERM_WRITE, + shortTermWrites, + MemoryUpdate::getPath, + MemoryUpdate::getValue, + settings); + return events; + } + + public static List buildReadEvents( + String eventKeyText, + List sensoryReads, + List shortTermReads, + MemoryEventSettings settings) { + List events = new ArrayList<>(2); + addFoldedEvent( + events, + eventKeyText, + MemoryOp.SENSORY_READ, + sensoryReads, + MemoryValueObservation::getPath, + MemoryValueObservation::getValue, + settings); + addFoldedEvent( + events, + eventKeyText, + MemoryOp.SHORT_TERM_READ, + shortTermReads, + MemoryValueObservation::getPath, + MemoryValueObservation::getValue, + settings); + return events; + } + + public static List buildLtmEvents( + String eventKeyText, List> records, MemoryEventSettings settings) { + Map> updates = new LinkedHashMap<>(); + Map> gets = new LinkedHashMap<>(); + Map> searches = new LinkedHashMap<>(); + Set clearedSets = new LinkedHashSet<>(); + + for (Map record : records) { + LtmOperation operation = parseLtmOperation(record.get("op")); + String memorySet = nonEmptyString(record.get("set")); + if (operation == null || memorySet == null) { + continue; + } + String memoryId = nonEmptyString(record.get("id")); + String query = nonEmptyString(record.get("query")); + Object value = record.get("value"); + + if (operation == LtmOperation.DELETE_SET) { + updates.remove(memorySet); + clearedSets.add(memorySet); + } else if ((operation == LtmOperation.ADD || operation == LtmOperation.UPDATE) + && memoryId != null) { + putNormalized(updates, memorySet, memoryId, value); + } else if (operation == LtmOperation.DELETE && memoryId != null) { + updates.computeIfAbsent(memorySet, ignored -> new LinkedHashMap<>()) + .put(memoryId, null); + } else if (operation == LtmOperation.GET && memoryId != null) { + putNormalized(gets, memorySet, memoryId, value); + } else if (operation == LtmOperation.SEARCH && query != null) { + putNormalized(searches, memorySet, query, value); + } + } + + List events = new ArrayList<>(3); + if (settings.generate(MemoryOp.LONG_TERM_UPDATE) + && (!updates.isEmpty() || !clearedSets.isEmpty())) { + addSafely( + events, + () -> + new LongTermUpdateEvent( + eventKeyText, + toObjectMap(updates), + new ArrayList<>(clearedSets)), + MemoryOp.LONG_TERM_UPDATE); + } + addNestedEvent(events, eventKeyText, MemoryOp.LONG_TERM_GET, gets, settings); + addNestedEvent(events, eventKeyText, MemoryOp.LONG_TERM_SEARCH, searches, settings); + return events; + } + + /** + * Parses and validates the private, versioned Python-to-Java LTM observation records. Malformed + * individual records are skipped; a malformed envelope yields no records. + */ + @SuppressWarnings("unchecked") + public static List> parseLtmObservationRecords(String json) { + if (json == null || json.isEmpty()) { + return Collections.emptyList(); + } + try { + List raw = MAPPER.readValue(json, List.class); + List> records = new ArrayList<>(raw.size()); + for (Object item : raw) { + if (!(item instanceof Map)) { + LOG.warn("Skipping malformed LTM observation record"); + continue; + } + Map record = (Map) item; + if (isValidLtmRecord(record)) { + records.add(record); + } + } + return records; + } catch (Exception | LinkageError e) { + LOG.warn("Dropping malformed LTM observation payload"); + return Collections.emptyList(); + } + } + + /** Strict JSON round-trip shared by runtime memory-observation producers. */ + @Nullable + public static Object normalizeValue(@Nullable Object value) throws IOException { + return MAPPER.readValue(MAPPER.writeValueAsBytes(value), Object.class); + } + + private static boolean isValidLtmRecord(Map record) { + Object version = record.get("version"); + LtmOperation operation = parseLtmOperation(record.get("op")); + String memorySet = nonEmptyString(record.get("set")); + if (!(version instanceof Integer) + || ((Integer) version) != LTM_RECORD_VERSION + || operation == null + || memorySet == null) { + LOG.warn("Skipping malformed LTM observation record"); + return false; + } + String memoryId = nonEmptyString(record.get("id")); + if ((operation == LtmOperation.ADD + || operation == LtmOperation.UPDATE + || operation == LtmOperation.DELETE + || operation == LtmOperation.GET) + && memoryId == null) { + LOG.warn("Skipping malformed LTM observation record for operation '{}'", operation); + return false; + } + if (operation == LtmOperation.SEARCH && nonEmptyString(record.get("query")) == null) { + LOG.warn("Skipping malformed LTM search observation record"); + return false; + } + return true; + } + + @Nullable + private static LtmOperation parseLtmOperation(Object value) { + if (!(value instanceof String)) { + return null; + } + try { + return LtmOperation.valueOf((String) value); + } catch (IllegalArgumentException e) { + LOG.warn("Skipping LTM observation record with unknown operation '{}'", value); + return null; + } + } + + @Nullable + private static String nonEmptyString(Object value) { + return value instanceof String && !((String) value).isEmpty() ? (String) value : null; + } + + private static void addFoldedEvent( + List events, + String eventKeyText, + MemoryOp operation, + List records, + Function pathExtractor, + Function valueExtractor, + MemoryEventSettings settings) { + if (records == null || records.isEmpty() || !settings.generate(operation)) { + return; + } + Map values = new LinkedHashMap<>(); + for (T record : records) { + String path = pathExtractor.apply(record); + Object normalized = normalizeSafely(valueExtractor.apply(record), operation); + if (normalized == INVALID_VALUE) { + // Last observation wins even when its value cannot be represented. Retaining an + // older value would falsely report stale memory state. + values.remove(path); + } else { + values.put(path, normalized); + } + } + if (!values.isEmpty()) { + addSafely(events, () -> operation.createEvent(eventKeyText, values), operation); + } + } + + private static void addNestedEvent( + List events, + String eventKeyText, + MemoryOp operation, + Map> values, + MemoryEventSettings settings) { + if (!values.isEmpty() && settings.generate(operation)) { + addSafely( + events, + () -> operation.createEvent(eventKeyText, toObjectMap(values)), + operation); + } + } + + private static void putNormalized( + Map> destination, + String memorySet, + @Nullable String itemKey, + @Nullable Object value) { + if (itemKey == null) { + return; + } + Object normalized = normalizeSafely(value, null); + if (normalized == INVALID_VALUE) { + removeNestedValue(destination, memorySet, itemKey); + return; + } + destination + .computeIfAbsent(memorySet, ignored -> new LinkedHashMap<>()) + .put(itemKey, normalized); + } + + private static void removeNestedValue( + Map> destination, String memorySet, String itemKey) { + Map values = destination.get(memorySet); + if (values == null) { + return; + } + values.remove(itemKey); + if (values.isEmpty()) { + destination.remove(memorySet); + } + } + + private static void addSafely( + List events, Supplier supplier, MemoryOp operation) { + try { + events.add(supplier.get()); + } catch (RuntimeException e) { + LOG.warn( + "Skipping framework memory observation event for operation '{}' because its value is not JSON-compatible ({})", + operation, + e.getClass().getSimpleName()); + } + } + + private static Object normalizeSafely(@Nullable Object value, @Nullable MemoryOp operation) { + try { + return normalizeValue(value); + } catch (Exception | LinkageError e) { + LOG.warn( + "Skipping framework memory observation value{} because it is not JSON-compatible ({})", + operation == null ? "" : " for operation '" + operation + "'", + e.getClass().getSimpleName()); + return INVALID_VALUE; + } + } + + private static Map toObjectMap(Map> nested) { + return new LinkedHashMap<>(nested); + } +} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryEventSettings.java b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryEventSettings.java new file mode 100644 index 0000000000..e83a15e8ad --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryEventSettings.java @@ -0,0 +1,107 @@ +/* + * 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.runtime.memory; + +import org.apache.flink.agents.api.configuration.ConfigOption; +import org.apache.flink.agents.api.configuration.MemoryEventOptions; +import org.apache.flink.agents.api.event.LongTermGetEvent; +import org.apache.flink.agents.api.event.LongTermSearchEvent; +import org.apache.flink.agents.api.event.LongTermUpdateEvent; +import org.apache.flink.agents.api.event.MemoryEvent; +import org.apache.flink.agents.api.event.SensoryReadEvent; +import org.apache.flink.agents.api.event.SensoryWriteEvent; +import org.apache.flink.agents.api.event.ShortTermReadEvent; +import org.apache.flink.agents.api.event.ShortTermWriteEvent; + +import java.util.EnumMap; +import java.util.Map; +import java.util.function.BiFunction; + +/** + * Resolved per-operation memory-event switches. Resolution per op: sub-key explicit → master switch + * explicit → per-op built-in default. + */ +public final class MemoryEventSettings { + + /** + * The seven observable memory operations: config option, built-in default, and event subclass. + */ + public enum MemoryOp { + SHORT_TERM_WRITE(MemoryEventOptions.SHORT_TERM_WRITE, true, ShortTermWriteEvent::new), + SHORT_TERM_READ(MemoryEventOptions.SHORT_TERM_READ, false, ShortTermReadEvent::new), + SENSORY_WRITE(MemoryEventOptions.SENSORY_WRITE, true, SensoryWriteEvent::new), + SENSORY_READ(MemoryEventOptions.SENSORY_READ, false, SensoryReadEvent::new), + LONG_TERM_UPDATE(MemoryEventOptions.LONG_TERM_UPDATE, true, LongTermUpdateEvent::new), + LONG_TERM_GET(MemoryEventOptions.LONG_TERM_GET, true, LongTermGetEvent::new), + LONG_TERM_SEARCH(MemoryEventOptions.LONG_TERM_SEARCH, true, LongTermSearchEvent::new); + + final ConfigOption option; + final boolean defaultEnabled; + final BiFunction, MemoryEvent> factory; + + MemoryOp( + ConfigOption option, + boolean defaultEnabled, + BiFunction, MemoryEvent> factory) { + this.option = option; + this.defaultEnabled = defaultEnabled; + this.factory = factory; + } + + /** Constructs this operation's event subclass with the given key and value. */ + public MemoryEvent createEvent(String key, Map value) { + return factory.apply(key, value); + } + } + + private final Map resolved; + private final boolean anyEnabled; + + private MemoryEventSettings(Map resolved) { + this.resolved = resolved; + this.anyEnabled = resolved.containsValue(true); + } + + /** Builds settings from the raw agent config map ({@code AgentConfiguration.getConfData()}). */ + public static MemoryEventSettings from(Map confData) { + Object master = confData.get(MemoryEventOptions.MEMORY_GENERATE_EVENT.getKey()); + Map resolved = new EnumMap<>(MemoryOp.class); + for (MemoryOp op : MemoryOp.values()) { + Object sub = confData.get(op.option.getKey()); + boolean enabled; + if (sub != null) { + enabled = Boolean.parseBoolean(sub.toString()); + } else if (master != null) { + enabled = Boolean.parseBoolean(master.toString()); + } else { + enabled = op.defaultEnabled; + } + resolved.put(op, enabled); + } + return new MemoryEventSettings(resolved); + } + + public boolean generate(MemoryOp op) { + return resolved.get(op); + } + + /** False only when every operation is disabled. */ + public boolean anyEnabled() { + return anyEnabled; + } +} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryObjectImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryObjectImpl.java index b9550a6760..80d49c548f 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryObjectImpl.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryObjectImpl.java @@ -21,6 +21,8 @@ import org.apache.flink.agents.api.context.MemoryRef; import org.apache.flink.agents.api.context.MemoryUpdate; +import javax.annotation.Nullable; + import java.io.Serializable; import java.util.ArrayList; import java.util.Collections; @@ -44,13 +46,14 @@ private enum ItemType { private final MemoryStore store; private final List memoryUpdates; + @Nullable private final List memoryReads; private final String prefix; private final Runnable mailboxThreadChecker; public MemoryObjectImpl( MemoryType type, MemoryStore store, String prefix, List memoryUpdates) throws Exception { - this(type, store, prefix, () -> {}, memoryUpdates); + this(type, store, prefix, () -> {}, memoryUpdates, null); } public MemoryObjectImpl( @@ -58,7 +61,8 @@ public MemoryObjectImpl( MemoryStore store, String prefix, Runnable mailboxThreadChecker, - List memoryUpdates) + List memoryUpdates, + @Nullable List memoryReads) throws Exception { this.type = type; this.store = store; @@ -68,6 +72,7 @@ public MemoryObjectImpl( store.put(ROOT_KEY, new MemoryItem()); } this.memoryUpdates = memoryUpdates; + this.memoryReads = memoryReads; } @Override @@ -75,7 +80,8 @@ public MemoryObject get(String path) throws Exception { mailboxThreadChecker.run(); String absPath = fullPath(path); if (store.contains(absPath)) { - return new MemoryObjectImpl(type, store, absPath, memoryUpdates); + return new MemoryObjectImpl( + type, store, absPath, mailboxThreadChecker, memoryUpdates, memoryReads); } return null; } @@ -142,7 +148,8 @@ public MemoryObject newObject(String path, boolean overwrite) throws Exception { parentItem.getSubKeys().add(parts[parts.length - 1]); store.put(parent, parentItem); - return new MemoryObjectImpl(type, store, absPath, memoryUpdates); + return new MemoryObjectImpl( + type, store, absPath, mailboxThreadChecker, memoryUpdates, memoryReads); } @Override @@ -176,6 +183,7 @@ public Map getFields() throws Exception { result.put(name, "NestedObject"); } else { result.put(name, memItem.getValue()); + recordValueRead(absPath, memItem.getValue()); } } return result; @@ -193,11 +201,18 @@ public Object getValue() throws Exception { mailboxThreadChecker.run(); MemoryItem memItem = store.get(prefix); if (memItem != null && memItem.getType() == ItemType.VALUE) { + recordValueRead(prefix, memItem.getValue()); return memItem.getValue(); } return null; } + private void recordValueRead(String path, Object value) { + if (memoryReads != null) { + memoryReads.add(new MemoryValueObservation(path, value)); + } + } + private String fullPath(String path) { return (prefix.isEmpty() ? path : prefix + SEPARATOR + path); } @@ -248,6 +263,10 @@ public ItemType getType() { return type; } + public boolean isValue() { + return type == ItemType.VALUE; + } + public Object getValue() { return value; } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryValueObservation.java b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryValueObservation.java new file mode 100644 index 0000000000..4e987d298a --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/MemoryValueObservation.java @@ -0,0 +1,37 @@ +/* + * 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.runtime.memory; + +/** One observed read of a real memory value. Object nodes are not read observations. */ +public final class MemoryValueObservation { + private final String path; + private final Object value; + + public MemoryValueObservation(String path, Object value) { + this.path = path; + this.value = value; + } + + public String getPath() { + return path; + } + + public Object getValue() { + return value; + } +} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java index bfdad4124f..27e74620a8 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java @@ -21,6 +21,7 @@ import org.apache.flink.agents.api.OutputEvent; import org.apache.flink.agents.api.agents.AgentExecutionOptions; import org.apache.flink.agents.api.context.MemoryUpdate; +import org.apache.flink.agents.api.event.AgentRunBeginEvent; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.PythonFunction; @@ -29,9 +30,12 @@ import org.apache.flink.agents.runtime.actionstate.ActionState; import org.apache.flink.agents.runtime.actionstate.ActionStateStore; import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory; +import org.apache.flink.agents.runtime.memory.MemoryEventBuilder; +import org.apache.flink.agents.runtime.memory.MemoryObjectImpl; import org.apache.flink.agents.runtime.metrics.BuiltInMetrics; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import org.apache.flink.agents.runtime.python.operator.PythonActionTask; +import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor; import org.apache.flink.agents.runtime.utils.EventUtil; import org.apache.flink.annotation.VisibleForTesting; import org.apache.flink.api.common.operators.MailboxExecutor; @@ -54,10 +58,14 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import javax.annotation.Nullable; + import java.lang.reflect.Field; import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; @@ -78,6 +86,7 @@ public class ActionExecutionOperator extends AbstractStreamOperator, BoundedOneInput { private static final long serialVersionUID = 1L; + private static final String AGENT_RUN_BEGIN_ACTION_NAME = "agent_run_begin_action"; private static final Logger LOG = LoggerFactory.getLogger(ActionExecutionOperator.class); @@ -118,17 +127,27 @@ public class ActionExecutionOperator extends AbstractStreamOperator(agentPlan, inputIsJava); this.durableExecManager = new DurableExecutionManager(actionStateStore); + this.agentRunBeginEventEnabled = + Boolean.TRUE.equals( + agentPlan.getConfig().get(AgentExecutionOptions.AGENT_RUN_BEGIN_EVENT)); OperatorUtils.setChainStrategy(this, ChainingStrategy.ALWAYS); } @@ -232,15 +251,20 @@ public void processElement(StreamRecord record) throws Exception { stateManager.addPendingInputEvent(inputEvent); } else { // Otherwise, the new event is processed immediately. - processEvent(getCurrentKey(), inputEvent); + processInputEvent(getCurrentKey(), inputEvent); } } + /** Resolves one context key for an input and reuses it for the entire agent run. */ + private void processInputEvent(Object key, Event inputEvent) throws Exception { + processEvent(key, resolveContextKey(key), inputEvent); + } + /** * Processes an incoming event for the given key and may submit a new mail * `tryProcessActionTaskForKey` to continue processing. */ - private void processEvent(Object key, Event event) throws Exception { + private void processEvent(Object key, String contextKey, Event event) throws Exception { eventRouter.notifyEventProcessed(event); boolean isInputEvent = EventUtil.isInputEvent(event); @@ -263,6 +287,7 @@ private void processEvent(Object key, Event event) throws Exception { // If the event is an InputEvent, we mark that the key is currently being processed. stateManager.addProcessingKey(key); stateManager.initOrIncSequenceNumber(); + tryEmitAgentRunBeginEvent(key, contextKey, event); } // We then obtain the triggered action and add ActionTasks to the waiting processing // queue. @@ -276,13 +301,59 @@ private void processEvent(Object key, Event event) throws Exception { if (isInputEvent) { // If the event is an InputEvent, we submit a new mail to try processing the actions. - mailboxExecutor.submit(() -> tryProcessActionTaskForKey(key), "process action task"); + mailboxExecutor.submit( + () -> tryProcessActionTaskForKey(key, contextKey), "process action task"); + } + } + + /** + * Attempts to emit an {@link AgentRunBeginEvent} for the input before any action triggered by + * that input executes. + */ + private void tryEmitAgentRunBeginEvent(Object key, String contextKey, Event inputEvent) + throws Exception { + if (!agentRunBeginEventEnabled) { + return; + } + Map stm = new LinkedHashMap<>(); + Iterable> entries = + stateManager.getShortTermMemState().entries(); + if (entries != null) { + for (Map.Entry entry : entries) { + MemoryObjectImpl.MemoryItem item = entry.getValue(); + if (item != null + && item.isValue() + && !MemoryObjectImpl.ROOT_KEY.equals(entry.getKey())) { + try { + stm.put(entry.getKey(), MemoryEventBuilder.normalizeValue(item.getValue())); + } catch (Exception | LinkageError e) { + LOG.warn( + "Skipping non-JSON-compatible STM value in AgentRunBeginEvent ({})", + e.getClass().getSimpleName()); + } + } + } } + final AgentRunBeginEvent beginEvent; + try { + beginEvent = new AgentRunBeginEvent(contextKey, stm); + } catch (RuntimeException | LinkageError e) { + LOG.warn( + "Skipping AgentRunBeginEvent because its value snapshot is not JSON-compatible ({})", + e.getClass().getSimpleName()); + return; + } + if (inputEvent.hasSourceTimestamp()) { + beginEvent.setSourceTimestamp(inputEvent.getSourceTimestamp()); + } + beginEvent.setUpstreamEventId(inputEvent.getId()); + beginEvent.setUpstreamActionName(AGENT_RUN_BEGIN_ACTION_NAME); + processEvent(key, contextKey, beginEvent); } - private void tryProcessActionTaskForKey(Object key) { + private void tryProcessActionTaskForKey(Object key, String contextKey) { try { - processActionTaskForKey(key); + processActionTaskForKey(key, contextKey); } catch (Throwable t) { // MailboxExecutor.submit() stores task failures in its Future. Catch Throwable and // rethrow via execute() so Errors fail the task instead of leaving the key in-flight. @@ -295,7 +366,7 @@ private void tryProcessActionTaskForKey(Object key) { } } - private void processActionTaskForKey(Object key) throws Exception { + private void processActionTaskForKey(Object key, String contextKey) throws Exception { // 1. Get an action task for the key. setCurrentKey(key); @@ -318,7 +389,7 @@ private void processActionTaskForKey(Object key) throws Exception { // 2. Invoke the action task. contextManager.createAndSetRunnerContext( actionTask, - key, + contextKey, agentPlan, resourceCache, metricGroup, @@ -373,10 +444,23 @@ private void processActionTaskForKey(Object key) throws Exception { durableExecManager.setupDurableExecutionContext( actionTask, actionState, sequenceNumber); - ActionTask.ActionTaskResult actionTaskResult = - actionTask.invoke( - getRuntimeContext().getUserCodeClassLoader(), - this.pythonBridge.getPythonActionExecutor()); + ActionTask.ActionTaskResult actionTaskResult; + try { + actionTaskResult = + actionTask.invoke( + getRuntimeContext().getUserCodeClassLoader(), + this.pythonBridge.getPythonActionExecutor()); + } catch (Throwable actionFailure) { + try { + actionTask.getRunnerContext().discardMemoryObservation(); + } catch (Throwable discardFailure) { + if (discardFailure != actionFailure) { + actionFailure.addSuppressed(discardFailure); + } + } + ExceptionUtils.rethrowException(actionFailure); + throw new AssertionError("Unreachable after rethrowing action failure"); + } // We remove the contexts from the map after the task is processed. They will be added // back later if the action task has a generated action task, meaning it is not @@ -398,7 +482,7 @@ private void processActionTaskForKey(Object key) throws Exception { } for (Event actionOutputEvent : outputEvents) { - processEvent(key, actionOutputEvent); + processEvent(key, contextKey, actionOutputEvent); } boolean currentInputEventFinished = false; @@ -445,12 +529,13 @@ private void processActionTaskForKey(Object key) throws Exception { eventRouter.processEligibleWatermarks(super::processWatermark); Event pendingInputEvent = stateManager.pollNextPendingInputEvent(); if (pendingInputEvent != null) { - processEvent(key, pendingInputEvent); + processInputEvent(key, pendingInputEvent); } } else if (stateManager.hasMoreActionTasks()) { // If the current key has additional action tasks remaining, we should submit a new mail // to continue processing them. - mailboxExecutor.submit(() -> tryProcessActionTaskForKey(key), "process action task"); + mailboxExecutor.submit( + () -> tryProcessActionTaskForKey(key, contextKey), "process action task"); } } @@ -555,6 +640,28 @@ private ActionTask createActionTask(Object key, Action action, Event event) { } } + /** Returns one textual context key for Java and PyFlink keyed streams. */ + private String resolveContextKey(Object key) { + PythonActionExecutor pythonActionExecutor = + pythonBridge == null ? null : pythonBridge.getPythonActionExecutor(); + return resolveContextKey(key, inputIsJava, pythonKeyIsPickled, pythonActionExecutor); + } + + @VisibleForTesting + static String resolveContextKey( + Object key, + boolean inputIsJava, + boolean pythonKeyIsPickled, + @Nullable PythonActionExecutor pythonActionExecutor) { + if (inputIsJava) { + return String.valueOf(key); + } + checkState( + pythonActionExecutor != null, + "PythonActionExecutor must be initialized for a PyFlink keyed stream"); + return pythonActionExecutor.resolveKeyText(key, pythonKeyIsPickled); + } + private void tryResumeProcessActionTasks() throws Exception { Iterable keys = stateManager.getProcessingKeys(); if (keys != null) { @@ -572,8 +679,9 @@ private void tryResumeProcessActionTasks() throws Exception { continue; } eventRouter.getKeySegmentQueue().addKeyToLastSegment(key); + String contextKey = resolveContextKey(key); mailboxExecutor.submit( - () -> tryProcessActionTaskForKey(key), "process action task"); + () -> tryProcessActionTaskForKey(key, contextKey), "process action task"); } stateManager.replaceProcessingKeys(new ArrayList<>(ownedKeys)); } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java index e6a0e21368..3eeec4bf1a 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorFactory.java @@ -26,6 +26,8 @@ import org.apache.flink.streaming.api.operators.StreamOperator; import org.apache.flink.streaming.api.operators.StreamOperatorParameters; +import static org.apache.flink.util.Preconditions.checkArgument; + /** Operator factory for {@link ActionExecutionOperator}. */ public class ActionExecutionOperatorFactory extends AbstractStreamOperatorFactory implements OneInputStreamOperatorFactory { @@ -34,17 +36,36 @@ public class ActionExecutionOperatorFactory extends AbstractStreamOpera private final Boolean inputIsJava; + private final boolean pythonKeyIsPickled; + private final ActionStateStore actionStateStore; public ActionExecutionOperatorFactory(AgentPlan agentPlan, Boolean inputIsJava) { - this(agentPlan, inputIsJava, null); + this(agentPlan, inputIsJava, false, null); + checkArgument( + Boolean.TRUE.equals(inputIsJava), + "PyFlink input must provide whether its logical key uses pickle serialization"); + } + + public ActionExecutionOperatorFactory( + AgentPlan agentPlan, Boolean inputIsJava, boolean pythonKeyIsPickled) { + this(agentPlan, inputIsJava, pythonKeyIsPickled, null); } @VisibleForTesting protected ActionExecutionOperatorFactory( AgentPlan agentPlan, Boolean inputIsJava, ActionStateStore actionStateStore) { + this(agentPlan, inputIsJava, false, actionStateStore); + } + + private ActionExecutionOperatorFactory( + AgentPlan agentPlan, + Boolean inputIsJava, + boolean pythonKeyIsPickled, + ActionStateStore actionStateStore) { this.agentPlan = agentPlan; this.inputIsJava = inputIsJava; + this.pythonKeyIsPickled = pythonKeyIsPickled; this.actionStateStore = actionStateStore; this.chainingStrategy = ChainingStrategy.ALWAYS; } @@ -56,6 +77,7 @@ public > T createStreamOperator( new ActionExecutionOperator<>( agentPlan, inputIsJava, + pythonKeyIsPickled, parameters.getProcessingTimeService(), parameters.getMailboxExecutor(), actionStateStore); diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java index 36d92ea434..e4972d5cb2 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java @@ -29,6 +29,7 @@ import java.util.List; import java.util.Objects; import java.util.Optional; +import java.util.UUID; /** * This class represents a task related to the execution of an action in {@link @@ -47,6 +48,10 @@ public abstract class ActionTask { protected final Object key; protected final Event event; protected final Action action; + + /** Stable identifier for observations produced by this logical action execution. */ + protected String observationId; + /** * Since RunnerContextImpl contains references to the Operator and state, it should not be * serialized and included in the state with ActionTask. Instead, we should check if a valid @@ -55,9 +60,14 @@ public abstract class ActionTask { protected transient RunnerContextImpl runnerContext; public ActionTask(Object key, Event event, Action action) { + this(key, event, action, UUID.randomUUID().toString()); + } + + protected ActionTask(Object key, Event event, Action action, String observationId) { this.key = key; this.event = event; this.action = action; + this.observationId = Objects.requireNonNull(observationId, "observationId"); } public RunnerContextImpl getRunnerContext() { @@ -72,6 +82,15 @@ public Object getKey() { return key; } + public String getObservationId() { + if (observationId == null) { + // Tasks restored from state written before observation IDs were introduced have no + // value for this field. Assign one before the task is invoked or checkpointed again. + observationId = UUID.randomUUID().toString(); + } + return observationId; + } + @Override public boolean equals(Object o) { if (this == o) return true; @@ -79,12 +98,13 @@ public boolean equals(Object o) { ActionTask other = (ActionTask) o; return Objects.equals(this.key, other.key) && Objects.equals(this.event, other.event) - && Objects.equals(this.action, other.action); + && Objects.equals(this.action, other.action) + && Objects.equals(this.getObservationId(), other.getObservationId()); } @Override public int hashCode() { - return Objects.hash(key, event, action); + return Objects.hash(key, event, action, getObservationId()); } /** Invokes the action task. */ diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java index eba25d5d45..9162a84485 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java @@ -17,6 +17,7 @@ */ package org.apache.flink.agents.runtime.operator; +import org.apache.flink.agents.api.event.MemoryEvent; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.JavaFunction; import org.apache.flink.agents.plan.PythonFunction; @@ -156,7 +157,7 @@ RunnerContextImpl createOrGetRunnerContext( * * * @param actionTask the task to be set up before execution. - * @param key the current Flink key. + * @param contextKey the textual key shared by LTM isolation and framework observation events. * @param agentPlan the agent plan. * @param resourceCache the resource cache. * @param metricGroup the agent metric group. @@ -169,7 +170,7 @@ RunnerContextImpl createOrGetRunnerContext( */ void createAndSetRunnerContext( ActionTask actionTask, - Object key, + String contextKey, AgentPlan agentPlan, ResourceCache resourceCache, FlinkAgentsMetricGroupImpl metricGroup, @@ -218,7 +219,11 @@ void createAndSetRunnerContext( } context.switchActionContext( - actionTask.action.getName(), memoryContext, String.valueOf(key.hashCode())); + actionTask.action.getName(), + memoryContext, + contextKey, + actionTask.getObservationId(), + MemoryEvent.isMemoryType(actionTask.event.getType())); if (context instanceof JavaRunnerContextImpl) { ContinuationContext continuationContext; diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java index 11724ce685..19bd677f9f 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java @@ -83,7 +83,9 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec if (finished) { return new ActionTaskResult( - true, runnerContext.drainEvents(event.getSourceTimestamp()), null); + true, + runnerContext.drainEventsAtActionFinish(event.getSourceTimestamp()), + null); } else { return new ActionTaskResult(false, Collections.emptyList(), this); } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java index cea9cce541..95e215d731 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java @@ -27,6 +27,7 @@ import org.apache.flink.agents.runtime.env.EmbeddedPythonEnvironment; import org.apache.flink.agents.runtime.env.PythonEnvironmentManager; import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory; +import org.apache.flink.agents.runtime.memory.MemoryEventSettings; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl; import org.apache.flink.agents.runtime.python.utils.JavaResourceAdapter; @@ -173,7 +174,7 @@ void open( initPythonActionExecutor(agentPlan, jobIdentifier); } if (mem0Configured) { - wireLongTermMemory(); + wireLongTermMemory(agentPlan); } initialized = true; } @@ -220,7 +221,7 @@ private boolean isMem0Configured(AgentPlan agentPlan) { * {@code create_flink_runner_context} already initialised via {@code _init_long_term_memory}) * and wrap it as a Java {@link Mem0LongTermMemory}. */ - private void wireLongTermMemory() { + private void wireLongTermMemory(AgentPlan agentPlan) { PyObject pyCtx = pythonActionExecutor.getPythonRunnerContext(); Object pyLtm = pythonInterpreter.invoke("python_java_utils.get_long_term_memory", pyCtx); if (pyLtm == null) { @@ -235,6 +236,13 @@ private void wireLongTermMemory() { LongTermMemoryOptions.Mem0.VECTOR_STORE.getKey())); } longTermMemory = new Mem0LongTermMemory(pythonResourceAdapter, (PyObject) pyLtm); + MemoryEventSettings settings = MemoryEventSettings.from(agentPlan.getConfigData()); + longTermMemory.configureObservation( + settings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_UPDATE), + settings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_GET), + settings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_SEARCH)); + // The Python runner context drains LTM observation records at action finish. + pythonRunnerContext.setLongTermMemory(longTermMemory); } private void initPythonActionExecutor(AgentPlan agentPlan, String jobIdentifier) diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java index ce43195a59..98a3c5af2b 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java @@ -39,6 +39,11 @@ public PythonActionTask(Object key, Event event, Action action) { checkState(action.getExec() instanceof PythonFunction); } + protected PythonActionTask(Object key, Event event, Action action, String observationId) { + super(key, event, action, observationId); + checkState(action.getExec() instanceof PythonFunction); + } + public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExecutor executor) throws Exception { LOG.debug( @@ -49,8 +54,7 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec runnerContext.checkNoPendingEvents(); String pythonAwaitableRef = - executor.executePythonFunction( - (PythonFunction) action.getExec(), event, key.hashCode()); + executor.executePythonFunction((PythonFunction) action.getExec(), event); // If a user-defined action uses an interface to submit asynchronous tasks, it will return a // Python coroutine (awaitable) object instance upon its first execution. Otherwise, it // means that no asynchronous tasks were submitted and the action has already completed. @@ -58,11 +62,12 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec // The Python action generates an awaitable. We need to execute it once, which will // submit an asynchronous task and return whether the action has been completed. ((PythonRunnerContextImpl) runnerContext).setPythonAwaitableRef(pythonAwaitableRef); - ActionTask tempGeneratedActionTask = new PythonGeneratorActionTask(key, event, action); + ActionTask tempGeneratedActionTask = + new PythonGeneratorActionTask(key, event, action, getObservationId()); tempGeneratedActionTask.setRunnerContext(runnerContext); return tempGeneratedActionTask.invoke(userCodeClassLoader, executor); } return new ActionTaskResult( - true, runnerContext.drainEvents(event.getSourceTimestamp()), null); + true, runnerContext.drainEventsAtActionFinish(event.getSourceTimestamp()), null); } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java index 35296c1cc5..a6d78eae14 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java @@ -26,8 +26,8 @@ /** An {@link ActionTask} wrapper a Python awaitable to represent a code block in Python action. */ public class PythonGeneratorActionTask extends PythonActionTask { - public PythonGeneratorActionTask(Object key, Event event, Action action) { - super(key, event, action); + public PythonGeneratorActionTask(Object key, Event event, Action action, String observationId) { + super(key, event, action, observationId); } @Override @@ -47,7 +47,8 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec "Python awaitable ref is null for action {} (likely restored from checkpoint), " + "re-executing from beginning.", action.getName()); - PythonActionTask freshTask = new PythonActionTask(key, event, action); + PythonActionTask freshTask = + new PythonActionTask(key, event, action, getObservationId()); freshTask.setRunnerContext(runnerContext); return freshTask.invoke(userCodeClassLoader, executor); } @@ -56,7 +57,9 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec ActionTask generatedActionTask = finished ? null : this; return new ActionTaskResult( finished, - runnerContext.drainEvents(event.getSourceTimestamp()), + finished + ? runnerContext.drainEventsAtActionFinish(event.getSourceTimestamp()) + : runnerContext.drainEvents(event.getSourceTimestamp()), generatedActionTask); } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java index 67c80f38db..d374027795 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java @@ -24,6 +24,7 @@ import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.PythonFunction; import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl; +import org.apache.flink.types.Row; import pemja.core.PythonInterpreter; import pemja.core.object.PyObject; @@ -44,9 +45,6 @@ public class PythonActionExecutor { private static final String CREATE_FLINK_RUNNER_CONTEXT = "flink_runner_context.create_flink_runner_context"; - private static final String FLINK_RUNNER_CONTEXT_SWITCH_ACTION_CONTEXT = - "flink_runner_context.flink_runner_context_switch_action_context"; - private static final String CLOSE_FLINK_RUNNER_CONTEXT = "flink_runner_context.close_flink_runner_context"; @@ -64,6 +62,10 @@ public class PythonActionExecutor { // =========== PYTHON AND JAVA OBJECT CONVERT =========== private static final String CONVERT_JSON_TO_PYTHON_EVENT = "python_java_utils.convert_json_to_python_event"; + private static final String CONVERT_TO_PYTHON_KEY_TEXT = + "python_java_utils.convert_to_python_key_text"; + private static final String PICKLED_KEY_SERIALIZATION = "pickled"; + private static final String EXPLICIT_KEY_SERIALIZATION = "explicit"; private static final String WRAP_TO_INPUT_EVENT = "python_java_utils.wrap_to_input_event"; private static final String GET_OUTPUT_FROM_OUTPUT_EVENT = "python_java_utils.get_output_from_output_event"; @@ -124,14 +126,10 @@ public void open() throws Exception { * @return The name of the Python awaitable variable. It may be null if the Python function does * not return a coroutine. */ - public String executePythonFunction(PythonFunction function, Event event, int hashOfKey) - throws Exception { + public String executePythonFunction(PythonFunction function, Event event) throws Exception { runnerContext.checkNoPendingEvents(); function.setInterpreter(interpreter); - interpreter.invoke( - FLINK_RUNNER_CONTEXT_SWITCH_ACTION_CONTEXT, pythonRunnerContext, hashOfKey); - String eventJson = new ObjectMapper().writeValueAsString(event); Object pythonEventObject = interpreter.invoke(CONVERT_JSON_TO_PYTHON_EVENT, eventJson); @@ -161,6 +159,22 @@ public Event wrapToInputEvent(Object eventData) throws IOException { return Event.fromJson((String) result); } + /** Resolves the textual logical key from PyFlink's keyed-stream representation. */ + public String resolveKeyText(Object flinkKey, boolean pythonKeyIsPickled) { + Object logicalKey = flinkKey; + if (flinkKey instanceof Row) { + logicalKey = ((Row) flinkKey).getField(0); + } + if (pythonKeyIsPickled || logicalKey instanceof byte[]) { + String keySerialization = + pythonKeyIsPickled ? PICKLED_KEY_SERIALIZATION : EXPLICIT_KEY_SERIALIZATION; + return (String) + interpreter.invoke( + CONVERT_TO_PYTHON_KEY_TEXT, (byte[]) logicalKey, keySerialization); + } + return String.valueOf(logicalKey); + } + public Object getOutputFromOutputEvent(String eventJson) { return interpreter.invoke(GET_OUTPUT_FROM_OUTPUT_EVENT, eventJson); } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java index ebfc15f468..5d53c18741 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/CompileUtilsTest.java @@ -19,11 +19,13 @@ import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorTest; +import org.apache.flink.api.common.typeinfo.Types; import org.apache.flink.api.java.functions.KeySelector; import org.apache.flink.streaming.api.datastream.DataStream; import org.apache.flink.streaming.api.datastream.DataStreamSource; import org.apache.flink.streaming.api.datastream.KeyedStream; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.typeinfo.python.PickledByteArrayTypeInfo; import org.apache.flink.util.CloseableIterator; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -96,6 +98,18 @@ void testJavaKeyedStreamConnectToAgent() throws Exception { checkResult(resultList); } + @Test + void detectsPickledAndExplicitByteArrayPythonKeyTypes() { + assertThat( + CompileUtils.isPickledPythonKeyType( + Types.ROW(PickledByteArrayTypeInfo.PICKLED_BYTE_ARRAY_TYPE_INFO))) + .isTrue(); + assertThat( + CompileUtils.isPickledPythonKeyType( + Types.ROW(Types.PRIMITIVE_ARRAY(Types.BYTE)))) + .isFalse(); + } + private static List getTestSequence() { List testSequence = new ArrayList<>(); for (int i = 0; i < TEST_SEQUENCE_REPEAT; i++) { diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java index f8dbc6fa8b..c35e1988b9 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java @@ -23,10 +23,13 @@ import org.apache.flink.agents.api.chat.messages.ChatMessage; import org.apache.flink.agents.api.chat.messages.MessageRole; import org.apache.flink.agents.api.context.MemoryUpdate; +import org.apache.flink.agents.api.event.AgentRunBeginEvent; import org.apache.flink.agents.api.event.ChatRequestEvent; import org.apache.flink.agents.api.event.ChatResponseEvent; import org.apache.flink.agents.api.event.ContextRetrievalRequestEvent; import org.apache.flink.agents.api.event.ContextRetrievalResponseEvent; +import org.apache.flink.agents.api.event.MemoryEvent; +import org.apache.flink.agents.api.event.ShortTermWriteEvent; import org.apache.flink.agents.api.event.ToolRequestEvent; import org.apache.flink.agents.api.event.ToolResponseEvent; import org.apache.flink.agents.api.tools.ToolResponse; @@ -37,6 +40,7 @@ import java.util.ArrayList; import java.util.Base64; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.UUID; @@ -431,6 +435,50 @@ public void testUnknownEnvelopeVersionRejected() throws Exception { assertThrows(RuntimeException.class, () -> ActionStateSerde.deserialize(patched)); } + @Test + public void testMemoryEventsSurviveActionStateRoundTrip() throws Exception { + ActionState state = new ActionState(new InputEvent("test input")); + Map value = new LinkedHashMap<>(); + value.put("bytes", new byte[] {1, 2, 3}); + value.put("pojo", new ObservationValuePojo("hello", 7)); + state.addEvent(new ShortTermWriteEvent("user-42", value)); + state.addEvent(new AgentRunBeginEvent("user-42", value)); + + MemoryEvent liveMemory = (MemoryEvent) state.getOutputEvents().get(0); + AgentRunBeginEvent liveRunBegin = (AgentRunBeginEvent) state.getOutputEvents().get(1); + + ActionState restored = ActionStateSerde.deserialize(ActionStateSerde.serialize(state)); + + assertEquals(2, restored.getOutputEvents().size()); + MemoryEvent memory = (MemoryEvent) restored.getOutputEvents().get(0); + AgentRunBeginEvent runBegin = (AgentRunBeginEvent) restored.getOutputEvents().get(1); + + assertEquals(liveMemory.getValue(), memory.getValue()); + assertEquals(liveRunBegin.getValue(), runBegin.getValue()); + assertEquals("AQID", memory.getValue().get("bytes")); + assertEquals(Map.of("name", "hello", "count", 7), memory.getValue().get("pojo")); + assertEquals("AQID", runBegin.getValue().get("bytes")); + assertEquals(Map.of("name", "hello", "count", 7), runBegin.getValue().get("pojo")); + } + + public static class ObservationValuePojo { + private final String name; + private final int count; + + public ObservationValuePojo(String name, int count) { + this.name = name; + this.count = count; + } + + public String getName() { + return name; + } + + public int getCount() { + return count; + } + } + /** Serializable POJO used to verify Kryo preserves user types across durable recovery. */ public static class MemoryValuePojo implements java.io.Serializable { private String name; diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestMemoryObservationFlush.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestMemoryObservationFlush.java new file mode 100644 index 0000000000..0bb01ad31f --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestMemoryObservationFlush.java @@ -0,0 +1,418 @@ +/* + * 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.runtime.context; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.event.LongTermUpdateEvent; +import org.apache.flink.agents.api.event.ShortTermWriteEvent; +import org.apache.flink.agents.api.memory.MemorySet; +import org.apache.flink.agents.api.memory.MemorySetItem; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.runtime.memory.CachedMemoryStore; +import org.apache.flink.agents.runtime.memory.ForTestMemoryMapState; +import org.apache.flink.agents.runtime.memory.InteranlBaseLongTermMemory; +import org.apache.flink.agents.runtime.memory.MemoryEventSettings; +import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; +import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups; +import org.junit.jupiter.api.Test; + +import javax.annotation.Nullable; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; + +/** Action-finish observation semantics at the real RunnerContext boundary. */ +class TestMemoryObservationFlush { + private RunnerContextImpl createContext(Map conf, boolean suppressed) + throws Exception { + return createContext(conf, suppressed, "user-42"); + } + + private RunnerContextImpl createContext( + Map conf, boolean suppressed, String contextKey) throws Exception { + return createContext(conf, suppressed, contextKey, null); + } + + private RunnerContextImpl createContext( + Map conf, + boolean suppressed, + String contextKey, + @Nullable InteranlBaseLongTermMemory ltm) + throws Exception { + AgentPlan plan = + new AgentPlan( + new HashMap<>(), + new HashMap<>(), + new org.apache.flink.agents.plan.AgentConfiguration(conf)); + RunnerContextImpl context = + new RunnerContextImpl( + new FlinkAgentsMetricGroupImpl( + UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup()), + () -> {}, + plan, + null, + "job"); + if (ltm != null) { + MemoryEventSettings settings = MemoryEventSettings.from(plan.getConfigData()); + ltm.configureObservation( + settings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_UPDATE), + settings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_GET), + settings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_SEARCH)); + context.setLongTermMemory(ltm); + } + context.switchActionContext( + "action", + new RunnerContextImpl.MemoryContext( + new CachedMemoryStore(new ForTestMemoryMapState<>()), + new CachedMemoryStore(new ForTestMemoryMapState<>())), + contextKey, + "observation-1", + suppressed); + return context; + } + + @Test + void finishReusesPersistenceWritesAndRecordsReadsFromGetFields() throws Exception { + Map conf = new HashMap<>(); + conf.put("memory.generate-event.short-term-read", true); + RunnerContextImpl context = createContext(conf, false); + context.getShortTermMemory().newObject("empty", false); + context.getShortTermMemory().set("user.tier", "gold"); + context.getShortTermMemory().get("user").getFields(); + + List events = context.drainEventsAtActionFinish(null); + + assertThat(events).hasSize(2); + assertThat(events.get(0).getType()).isEqualTo(ShortTermWriteEvent.EVENT_TYPE); + assertThat(((ShortTermWriteEvent) events.get(0)).getValue()) + .containsEntry("empty", null) + .containsEntry("user.tier", "gold"); + assertThat(events.get(1).getAttr("value").toString()).contains("user.tier=gold"); + assertReadObservationListsEmpty(context.getMemoryContext()); + } + + @Test + void invalidKryoOnlyOrNonFiniteValuesNeverFailActionFinish() throws Exception { + RunnerContextImpl context = createContext(new HashMap<>(), false); + context.getShortTermMemory().set("valid", "kept"); + context.getShortTermMemory().set("nan", Double.NaN); + context.getShortTermMemory().set("pojo", new KryoOnlyValue()); + + assertThatCode(() -> context.drainEventsAtActionFinish(null)).doesNotThrowAnyException(); + assertThat(context.getShortTermMemoryUpdates()).hasSize(3); + } + + @Test + void unfinishedDrainDefersObservationsUntilFinish() throws Exception { + RunnerContextImpl context = createContext(new HashMap<>(), false); + context.getShortTermMemory().set("first", 1); + + assertThat(context.drainEvents(null)).isEmpty(); + context.getShortTermMemory().set("second", 2); + + List finished = context.drainEventsAtActionFinish(null); + assertThat(finished).singleElement(); + assertThat(((ShortTermWriteEvent) finished.get(0)).getValue()) + .containsEntry("first", 1) + .containsEntry("second", 2); + assertThat(context.getShortTermMemoryUpdates()).hasSize(2); + } + + @Test + void userAndMemoryEventsCoexistAtFinish() throws Exception { + RunnerContextImpl context = createContext(new HashMap<>(), false); + context.sendEvent(new Event("user-event", Map.of("answer", 42))); + context.getShortTermMemory().set("memory", "value"); + + assertThat(context.drainEventsAtActionFinish(null)) + .extracting(Event::getType) + .containsExactly("user-event", ShortTermWriteEvent.EVENT_TYPE); + } + + @Test + void disabledLtmOperationsDoNotCrossBridgeAtFinishOrDiscard() throws Exception { + StubLtm ltm = new StubLtm(); + RunnerContextImpl context = + createContext(Map.of("memory.generate-event", false), false, "user-42", ltm); + + assertThat(context.drainEventsAtActionFinish(null)).isEmpty(); + context.discardMemoryObservation(); + + assertThat(ltm.drainCallCount).isZero(); + assertThat(ltm.updateEnabled).isFalse(); + assertThat(ltm.getEnabled).isFalse(); + assertThat(ltm.searchEnabled).isFalse(); + } + + @Test + void individualLtmOperationFlagsAreForwardedExactly() throws Exception { + Map conf = new HashMap<>(); + conf.put("memory.generate-event.long-term-update", false); + conf.put("memory.generate-event.long-term-get", true); + conf.put("memory.generate-event.long-term-search", false); + StubLtm ltm = new StubLtm(); + RunnerContextImpl context = createContext(conf, false, "user-42", ltm); + + assertThat(context.drainEventsAtActionFinish(null)).isEmpty(); + + assertThat(ltm.updateEnabled).isFalse(); + assertThat(ltm.getEnabled).isTrue(); + assertThat(ltm.searchEnabled).isFalse(); + assertThat(ltm.drainCallCount).isEqualTo(1); + } + + @Test + void suppressedActionDoesNotCrossLtmBridge() throws Exception { + StubLtm suppressedLtm = new StubLtm(); + RunnerContextImpl suppressed = + createContext(new HashMap<>(), true, "user-42", suppressedLtm); + suppressed.drainEventsAtActionFinish(null); + suppressed.discardMemoryObservation(); + + assertThat(suppressedLtm.drainCallCount).isZero(); + assertThat(suppressedLtm.observationSuppressed).isTrue(); + } + + @Test + void textualContextKeyRecordsMemoryObservationsAndPersistsUpdates() throws Exception { + Map conf = new HashMap<>(); + conf.put("memory.generate-event.sensory-read", true); + conf.put("memory.generate-event.short-term-read", true); + RunnerContextImpl context = createContext(conf, false, "42"); + + context.getSensoryMemory().set("sensory", "value"); + context.getSensoryMemory().get("sensory").getValue(); + context.getShortTermMemory().set("short-term", "value"); + context.getShortTermMemory().get("short-term").getValue(); + + assertThat(context.getSensoryMemoryUpdates()).hasSize(1); + assertThat(context.getShortTermMemoryUpdates()).hasSize(1); + assertThat(context.drainEventsAtActionFinish(null)) + .hasSize(4) + .allSatisfy(event -> assertThat(event.getAttr("key")).isEqualTo("42")); + assertReadObservationListsEmpty(context.getMemoryContext()); + } + + @Test + void observationConfigurationIsNotRepeatedAcrossActionSwitches() throws Exception { + StubLtm ltm = new StubLtm(); + RunnerContextImpl context = createContext(new HashMap<>(), false, "user-42", ltm); + + context.switchActionContext( + "next-action", + new RunnerContextImpl.MemoryContext( + new CachedMemoryStore(new ForTestMemoryMapState<>()), + new CachedMemoryStore(new ForTestMemoryMapState<>())), + "user-43", + "observation-2", + true); + + assertThat(ltm.configureCallCount).isEqualTo(1); + assertThat(ltm.switchCallCount).isEqualTo(2); + assertThat(ltm.partitionKey).isEqualTo("user-43"); + assertThat(ltm.observationId).isEqualTo("observation-2"); + assertThat(ltm.observationSuppressed).isTrue(); + } + + @Test + void validVersionedLtmRecordUsesStructuredSetMap() throws Exception { + StubLtm ltm = new StubLtm(); + ltm.payload = + "[{\"version\":1,\"op\":\"ADD\",\"set\":\"a.b\",\"id\":\"m.1\",\"value\":\"v\"}]"; + RunnerContextImpl context = createContext(new HashMap<>(), false, "user-42", ltm); + + LongTermUpdateEvent event = + (LongTermUpdateEvent) context.drainEventsAtActionFinish(null).get(0); + assertThat(event.getValue()).containsEntry("a.b", Map.of("m.1", "v")); + } + + @Test + void interleavedSameKeyActionsKeepLtmEventsWithTheirOwningExecution() throws Exception { + ActionScopedStubLtm ltm = new ActionScopedStubLtm(); + RunnerContextImpl context = createContext(new HashMap<>(), false, "user-42", ltm); + + // Action A records an LTM update and suspends. + ltm.record("user-42", "observation-1", "a", "from-a"); + + // Action B starts on the same key and finishes before A's continuation. + context.switchActionContext( + "action-b", + new RunnerContextImpl.MemoryContext( + new CachedMemoryStore(new ForTestMemoryMapState<>()), + new CachedMemoryStore(new ForTestMemoryMapState<>())), + "user-42", + "observation-2", + false); + ltm.record("user-42", "observation-2", "b", "from-b"); + + LongTermUpdateEvent bEvent = + (LongTermUpdateEvent) context.drainEventsAtActionFinish(null).get(0); + assertThat(bEvent.getValue()).containsExactly(Map.entry("test", Map.of("b", "from-b"))); + + // A resumes and fails: discarding A must not emit its record or affect B's event. + context.switchActionContext( + "action-a", + new RunnerContextImpl.MemoryContext( + new CachedMemoryStore(new ForTestMemoryMapState<>()), + new CachedMemoryStore(new ForTestMemoryMapState<>())), + "user-42", + "observation-1", + false); + context.discardMemoryObservation(); + assertThat(context.drainEventsAtActionFinish(null)).isEmpty(); + assertThat(ltm.pendingRecords).isEmpty(); + } + + @Test + void malformedPayloadAndLinkageErrorRemainBestEffort() throws Exception { + StubLtm malformedLtm = new StubLtm(); + malformedLtm.payload = "not-json"; + RunnerContextImpl malformed = + createContext(new HashMap<>(), false, "user-42", malformedLtm); + assertThatCode(() -> malformed.drainEventsAtActionFinish(null)).doesNotThrowAnyException(); + + StubLtm linkageLtm = new StubLtm(); + linkageLtm.drainLinkageError = new NoClassDefFoundError("missing-observation-codec"); + RunnerContextImpl linkage = createContext(new HashMap<>(), false, "user-42", linkageLtm); + + assertThatCode(() -> linkage.drainEventsAtActionFinish(null)).doesNotThrowAnyException(); + assertThat(linkageLtm.drainCallCount).isEqualTo(1); + } + + private static void assertReadObservationListsEmpty( + RunnerContextImpl.MemoryContext memoryContext) { + assertThat(memoryContext.getSensoryMemoryReads()).isEmpty(); + assertThat(memoryContext.getShortTermMemoryReads()).isEmpty(); + } + + private static final class KryoOnlyValue { + private final Object unsupported = new Object(); + } + + private static class StubLtm implements InteranlBaseLongTermMemory { + private String payload = "[]"; + private boolean updateEnabled; + private boolean getEnabled; + private boolean searchEnabled; + private String partitionKey; + private String observationId; + private boolean observationSuppressed; + private int configureCallCount; + private int switchCallCount; + private int drainCallCount; + private LinkageError drainLinkageError; + + @Override + public void configureObservation( + boolean updateObservationEnabled, + boolean getObservationEnabled, + boolean searchObservationEnabled) { + this.updateEnabled = updateObservationEnabled; + this.getEnabled = getObservationEnabled; + this.searchEnabled = searchObservationEnabled; + configureCallCount++; + } + + @Override + public void switchContext( + String partitionKey, String observationId, boolean observationSuppressed) { + this.partitionKey = partitionKey; + this.observationId = observationId; + this.observationSuppressed = observationSuppressed; + switchCallCount++; + } + + @Override + public String drainObservationRecordsJson(String partitionKey, String observationId) { + drainCallCount++; + if (drainLinkageError != null) { + throw drainLinkageError; + } + String result = payload; + payload = "[]"; + return result; + } + + @Override + public MemorySet getMemorySet(String name) { + return null; + } + + @Override + public boolean deleteMemorySet(String name) { + return false; + } + + @Override + public List add( + MemorySet memorySet, + List memoryItems, + @Nullable List> metadatas) { + return List.of(); + } + + @Override + public List get( + MemorySet memorySet, + @Nullable List ids, + @Nullable Map filters, + @Nullable Integer limit) { + return List.of(); + } + + @Override + public void delete(MemorySet memorySet, @Nullable List ids) {} + + @Override + public List search( + MemorySet memorySet, + String query, + int limit, + @Nullable Map filters, + Map extraArgs) { + return List.of(); + } + + @Override + public void close() {} + } + + private static final class ActionScopedStubLtm extends StubLtm { + private final Map pendingRecords = new HashMap<>(); + + private void record(String partitionKey, String observationId, String id, String value) { + pendingRecords.put( + partitionKey + "\u0000" + observationId, + "[{\"version\":1,\"op\":\"ADD\",\"set\":\"test\",\"id\":\"" + + id + + "\",\"value\":\"" + + value + + "\"}]"); + } + + @Override + public String drainObservationRecordsJson(String partitionKey, String observationId) { + String payload = pendingRecords.remove(partitionKey + "\u0000" + observationId); + return payload == null ? "[]" : payload; + } + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/ForTestMemoryMapState.java b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/ForTestMemoryMapState.java index 916d68cd6b..1abec5d378 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/ForTestMemoryMapState.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/ForTestMemoryMapState.java @@ -24,62 +24,62 @@ import java.util.Map; /** Simple, non-serialized HashMap implementation. */ -class ForTestMemoryMapState implements MapState { +public class ForTestMemoryMapState implements MapState { - private final Map fortest = new HashMap<>(); + private final Map delegate = new HashMap<>(); @Override public V get(String key) { - return fortest.get(key); + return delegate.get(key); } @Override public void put(String key, V value) { - fortest.put(key, value); + delegate.put(key, value); } @Override public void putAll(Map map) { - fortest.putAll(map); + delegate.putAll(map); } @Override public void remove(String key) { - fortest.remove(key); + delegate.remove(key); } @Override public boolean contains(String key) { - return fortest.containsKey(key); + return delegate.containsKey(key); } @Override public Iterable> entries() { - return fortest.entrySet(); + return delegate.entrySet(); } @Override public Iterable keys() { - return fortest.keySet(); + return delegate.keySet(); } @Override public Iterable values() { - return fortest.values(); + return delegate.values(); } @Override public Iterator> iterator() { - return fortest.entrySet().iterator(); + return delegate.entrySet().iterator(); } @Override public boolean isEmpty() { - return fortest.isEmpty(); + return delegate.isEmpty(); } @Override public void clear() { - fortest.clear(); + delegate.clear(); } } 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 eee167bdf9..f2575be45e 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 @@ -23,6 +23,7 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.MockitoAnnotations; import pemja.core.object.PyObject; @@ -32,7 +33,6 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -60,6 +60,13 @@ void tearDown() throws Exception { } } + @SuppressWarnings({"rawtypes", "unchecked"}) + private Map captureKwargs(String methodName) { + ArgumentCaptor captor = ArgumentCaptor.forClass(Map.class); + verify(mockAdapter).callMethod(eq(mockPyMem0), eq(methodName), captor.capture()); + return captor.getValue(); + } + @Test void testGetMemorySetIsPureFactoryAndBindsLtm() throws Exception { MemorySet ms = ltm.getMemorySet("notes"); @@ -96,18 +103,9 @@ void testAddForwardsKwargsAndReturnsIds() throws Exception { ltm.add(ms, List.of("hello", "world"), List.of(Map.of("k", "v"), Map.of())); assertThat(ids).containsExactly("a", "b"); - verify(mockAdapter) - .callMethod( - eq(mockPyMem0), - eq("add"), - argThat( - kwargs -> { - assertThat(kwargs) - .containsKeys( - "memory_set", "memory_items", "metadatas"); - assertThat(kwargs.get("memory_set")).isEqualTo(mockPyMemorySet); - return true; - })); + assertThat(captureKwargs("add")) + .containsKeys("memory_set", "memory_items", "metadatas") + .containsEntry("memory_set", mockPyMemorySet); } @Test @@ -119,15 +117,7 @@ void testGetOmitsNullOptionalKwargs() throws Exception { ltm.get(ms, null, null, null); - verify(mockAdapter) - .callMethod( - eq(mockPyMem0), - eq("get"), - argThat( - kwargs -> { - assertThat(kwargs).containsOnlyKeys("memory_set"); - return true; - })); + assertThat(captureKwargs("get")).containsOnlyKeys("memory_set"); } @Test @@ -153,16 +143,9 @@ void testGetWithIdsAndFiltersConvertsItems() throws Exception { assertThat(item.getAdditionalMetadata()).containsEntry("k", "v"); assertThat(item.getCreatedAt()).isNull(); - verify(mockAdapter) - .callMethod( - eq(mockPyMem0), - eq("get"), - argThat( - kwargs -> { - assertThat(kwargs).containsKeys("ids", "filters", "limit"); - assertThat(kwargs.get("limit")).isEqualTo(50); - return true; - })); + assertThat(captureKwargs("get")) + .containsKeys("ids", "filters", "limit") + .containsEntry("limit", 50); } @Test @@ -171,15 +154,7 @@ void testDeleteForwardsIds() throws Exception { ltm.delete(ms, List.of("id1", "id2")); - verify(mockAdapter) - .callMethod( - eq(mockPyMem0), - eq("delete"), - argThat( - kwargs -> { - assertThat(kwargs).containsKeys("memory_set", "ids"); - return true; - })); + assertThat(captureKwargs("delete")).containsKeys("memory_set", "ids"); } @Test @@ -188,15 +163,7 @@ void testDeleteWithoutIdsOmitsKwarg() throws Exception { ltm.delete(ms, null); - verify(mockAdapter) - .callMethod( - eq(mockPyMem0), - eq("delete"), - argThat( - kwargs -> { - assertThat(kwargs).containsOnlyKeys("memory_set"); - return true; - })); + assertThat(captureKwargs("delete")).containsOnlyKeys("memory_set"); } @Test @@ -208,33 +175,49 @@ void testSearchForwardsKwargs() throws Exception { ltm.search(ms, "hi", 10, Map.of("user_id", "u1"), Map.of("threshold", 0.7)); - verify(mockAdapter) - .callMethod( - eq(mockPyMem0), - eq("search"), - argThat( - kwargs -> { - assertThat(kwargs) - .containsKeys( - "memory_set", - "query", - "limit", - "filters", - "threshold"); - assertThat(kwargs.get("query")).isEqualTo("hi"); - assertThat(kwargs.get("limit")).isEqualTo(10); - assertThat(kwargs.get("threshold")).isEqualTo(0.7); - return true; - })); + assertThat(captureKwargs("search")) + .containsKeys("memory_set", "query", "limit", "filters", "threshold") + .containsEntry("query", "hi") + .containsEntry("limit", 10) + .containsEntry("threshold", 0.7); } @Test void testSwitchContextAndCloseForward() { - ltm.switchContext("k1"); + ltm.configureObservation(true, false, true); + ltm.switchContext("k1", "observation-1", true); + ltm.drainObservationRecordsJson("k1", "observation-1"); ltm.close(); verify(mockAdapter) - .callMethod(eq(mockPyMem0), eq("switch_context"), eq(Map.of("key", "k1"))); + .callMethod( + eq(mockPyMem0), + eq("configure_observation"), + eq( + Map.of( + "update_observation_enabled", + true, + "get_observation_enabled", + false, + "search_observation_enabled", + true))); + verify(mockAdapter) + .callMethod( + eq(mockPyMem0), + eq("switch_context"), + eq( + Map.of( + "key", + "k1", + "observation_id", + "observation-1", + "observation_suppressed", + true))); + verify(mockAdapter) + .callMethod( + eq(mockPyMem0), + eq("drain_ltm_observation_records"), + eq(Map.of("key", "k1", "observation_id", "observation-1"))); verify(mockAdapter).callMethod(eq(mockPyMem0), eq("close"), eq(Map.of())); } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/ShortTermMemoryTTLIntegrationTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/ShortTermMemoryTTLIntegrationTest.java index 97e4099e20..54a42e0c81 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/ShortTermMemoryTTLIntegrationTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/ShortTermMemoryTTLIntegrationTest.java @@ -45,13 +45,21 @@ class ShortTermMemoryTTLIntegrationTest { private static final class TestInput { public String eventKey; + public long delayBeforeMs; public long sleepMs; + public boolean write; private TestInput() {} private TestInput(String eventKey, long sleepMs) { + this(eventKey, 0L, sleepMs, true); + } + + private TestInput(String eventKey, long delayBeforeMs, long sleepMs, boolean write) { this.eventKey = eventKey; + this.delayBeforeMs = delayBeforeMs; this.sleepMs = sleepMs; + this.write = write; } } @@ -63,6 +71,7 @@ public static void input(org.apache.flink.agents.api.Event event, RunnerContext InputEvent inputEvent = (InputEvent) event; TestInput input = (TestInput) inputEvent.getInput(); + Thread.sleep(input.delayBeforeMs); MemoryObject shortTermMemory = ctx.getShortTermMemory(); MemoryObject memoryObject = shortTermMemory.get(input.eventKey); @@ -77,7 +86,9 @@ public static void input(org.apache.flink.agents.api.Event event, RunnerContext } } - shortTermMemory.set(input.eventKey, currentCount + 1); + if (input.write) { + shortTermMemory.set(input.eventKey, currentCount + 1); + } Thread.sleep(input.sleepMs); ctx.sendEvent( new OutputEvent( @@ -120,6 +131,49 @@ void testTTLConfigurationAppliedWithDefaultUpdateTypeAndVisibility() throws Exce assertEquals(List.of("event1|NEW", "event2|NEW", "event1|NEW"), results); } + @Test + void testDefaultUpdateTypeRefreshesTtlOnRead() throws Exception { + List defaultResults = runReadRefreshScenario(null); + List onCreateAndWriteResults = + runReadRefreshScenario(ShortTermMemoryTtlUpdate.ON_CREATE_AND_WRITE); + + assertEquals(List.of("event1|NEW", "event1|EXISTING", "event1|EXISTING"), defaultResults); + assertEquals( + List.of("event1|NEW", "event1|EXISTING", "event1|NEW"), onCreateAndWriteResults); + } + + private static List runReadRefreshScenario(ShortTermMemoryTtlUpdate updateType) + throws Exception { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + + AgentsExecutionEnvironment agentEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + AgentConfiguration agentsConfig = (AgentConfiguration) agentEnv.getConfig(); + agentsConfig.set(AgentExecutionOptions.SHORT_TERM_MEMORY_STATE_TTL_MS, 3000L); + if (updateType != null) { + agentsConfig.set( + AgentExecutionOptions.SHORT_TERM_MEMORY_STATE_TTL_UPDATE_TYPE, updateType); + } + + List testData = new ArrayList<>(); + // Reads occur near t=2s and t=4s. With a 3s TTL, the last read is after the + // original expiry but before the first read's refreshed expiry near t=5s. + testData.add(new TestInput("event1", 0L, 0L, true)); + testData.add(new TestInput("event1", 2000L, 0L, false)); + testData.add(new TestInput("event1", 2000L, 0L, false)); + + DataStream inputStream = env.fromCollection(testData); + DataStream outputStream = + agentEnv.fromDataStream(inputStream, x -> MEMORY_KEY) + .apply(new TTLTestAgent()) + .toDataStream(); + + List results = new ArrayList<>(); + outputStream.map(Object::toString).executeAndCollect().forEachRemaining(results::add); + return results; + } + private static List runScenario( long ttlMs, long sleepMs, boolean configureTtlMs, boolean configureTtlOptions) throws Exception { diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/TestMemoryEventBuilder.java b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/TestMemoryEventBuilder.java new file mode 100644 index 0000000000..1139d5c4d8 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/TestMemoryEventBuilder.java @@ -0,0 +1,155 @@ +/* + * 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.runtime.memory; + +import org.apache.flink.agents.api.context.MemoryUpdate; +import org.apache.flink.agents.api.event.LongTermGetEvent; +import org.apache.flink.agents.api.event.LongTermSearchEvent; +import org.apache.flink.agents.api.event.LongTermUpdateEvent; +import org.apache.flink.agents.api.event.MemoryEvent; +import org.apache.flink.agents.api.event.ShortTermWriteEvent; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Memory-update folding and versioned long-term-memory observation rules. */ +class TestMemoryEventBuilder { + private static MemoryEventSettings allOn() { + return MemoryEventSettings.from(Map.of("memory.generate-event", true)); + } + + @Test + void writesFoldLastValueAndKeepNullMarkers() { + List events = + MemoryEventBuilder.buildWriteEvents( + "user-42", + List.of(), + List.of( + new MemoryUpdate("user.tier", "silver"), + new MemoryUpdate("user.tier", "gold"), + new MemoryUpdate("nullable", null)), + allOn()); + + assertThat(events).hasSize(1); + assertThat(events.get(0).getType()).isEqualTo(ShortTermWriteEvent.EVENT_TYPE); + assertThat(events.get(0).getValue()) + .containsEntry("user.tier", "gold") + .containsEntry("nullable", null); + } + + @Test + void invalidValueIsSkippedWithoutDroppingValidValue() { + List events = + MemoryEventBuilder.buildWriteEvents( + "k", + List.of(), + List.of(new MemoryUpdate("bad", Double.NaN), new MemoryUpdate("good", "v")), + allOn()); + + assertThat(events).hasSize(1); + assertThat(events.get(0).getValue()).containsOnlyKeys("good"); + } + + @Test + void invalidLastShortTermValueRemovesEarlierFoldedValue() { + List events = + MemoryEventBuilder.buildWriteEvents( + "k", + List.of(), + List.of( + new MemoryUpdate("stale", "old"), + new MemoryUpdate("kept", "current"), + new MemoryUpdate("stale", Double.NaN)), + allOn()); + + assertThat(events).hasSize(1); + assertThat(events.get(0).getValue()) + .containsEntry("kept", "current") + .doesNotContainKey("stale"); + } + + @Test + void ltmUsesStructuredSetIdentityAndExplicitSetClear() { + List> records = + MemoryEventBuilder.parseLtmObservationRecords( + "[" + + "{\"version\":1,\"op\":\"ADD\",\"set\":\"a\",\"id\":\"b.m1\",\"value\":\"old\"}," + + "{\"version\":1,\"op\":\"ADD\",\"set\":\"a.b\",\"id\":\"m1\",\"value\":\"kept\"}," + + "{\"version\":1,\"op\":\"DELETE_SET\",\"set\":\"a\"}," + + "{\"version\":1,\"op\":\"UPDATE\",\"set\":\"a\",\"id\":\"m2\",\"value\":\"new\"}," + + "{\"version\":1,\"op\":\"GET\",\"set\":\"a.b\",\"id\":\"m1\",\"value\":\"kept\"}," + + "{\"version\":1,\"op\":\"SEARCH\",\"set\":\"a.b\",\"query\":\"q\",\"value\":[{\"id\":\"m1\"}]}" + + "]"); + + List events = MemoryEventBuilder.buildLtmEvents("k", records, allOn()); + + assertThat(events).hasSize(3); + LongTermUpdateEvent update = (LongTermUpdateEvent) events.get(0); + assertThat(update.getClearedSets()).containsExactly("a"); + assertThat(update.getValue()) + .containsEntry("a", Map.of("m2", "new")) + .containsEntry("a.b", Map.of("m1", "kept")); + assertThat(events.get(1).getType()).isEqualTo(LongTermGetEvent.EVENT_TYPE); + assertThat(events.get(1).getValue()).containsEntry("a.b", Map.of("m1", "kept")); + assertThat(events.get(2).getType()).isEqualTo(LongTermSearchEvent.EVENT_TYPE); + assertThat(events.get(2).getValue()).containsKey("a.b"); + } + + @Test + void parserRejectsUnknownOrUnversionedRecords() { + assertThat( + MemoryEventBuilder.parseLtmObservationRecords( + "[{\"op\":\"ADD\",\"set\":\"s\",\"id\":\"i\"}," + + "{\"version\":1,\"op\":\"UNKNOWN\",\"set\":\"s\"}]")) + .isEmpty(); + } + + @Test + @SuppressWarnings("unchecked") + void secondaryJavaNormalizationFailureRemovesEarlierLtmValue() { + List> records = + MemoryEventBuilder.parseLtmObservationRecords( + "[" + + "{\"version\":1,\"op\":\"UPDATE\",\"set\":\"s\",\"id\":\"stale\",\"value\":\"old\"}," + + "{\"version\":1,\"op\":\"UPDATE\",\"set\":\"s\",\"id\":\"stale\",\"value\":{}}," + + "{\"version\":1,\"op\":\"UPDATE\",\"set\":\"s\",\"id\":\"kept\",\"value\":\"new\"}" + + "]"); + ((Map) records.get(1).get("value")).put("invalid", Double.NaN); + + LongTermUpdateEvent event = + (LongTermUpdateEvent) + MemoryEventBuilder.buildLtmEvents("k", records, allOn()).get(0); + + assertThat(event.getValue()).containsEntry("s", Map.of("kept", "new")); + } + + @Test + void parserRequiresExactVersionAndNonEmptyIdentityFields() { + assertThat( + MemoryEventBuilder.parseLtmObservationRecords( + "[" + + "{\"version\":1.5,\"op\":\"ADD\",\"set\":\"s\",\"id\":\"i\"}," + + "{\"version\":1,\"op\":\"GET\",\"set\":\"s\",\"id\":\"\"}," + + "{\"version\":1,\"op\":\"SEARCH\",\"set\":\"s\",\"query\":\"\"}" + + "]")) + .isEmpty(); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/TestMemoryEventSettings.java b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/TestMemoryEventSettings.java new file mode 100644 index 0000000000..48acb524e5 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/TestMemoryEventSettings.java @@ -0,0 +1,87 @@ +/* + * 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.runtime.memory; + +import org.junit.jupiter.api.Test; + +import java.util.Map; + +import static org.apache.flink.agents.runtime.memory.MemoryEventSettings.MemoryOp; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Three-level fallback: sub-key explicit → master explicit → per-op default. */ +public class TestMemoryEventSettings { + + @Test + void testPerOpDefaultsWhenNothingConfigured() { + MemoryEventSettings settings = MemoryEventSettings.from(Map.of()); + assertTrue(settings.generate(MemoryOp.SHORT_TERM_WRITE)); + assertTrue(settings.generate(MemoryOp.SENSORY_WRITE)); + assertTrue(settings.generate(MemoryOp.LONG_TERM_UPDATE)); + assertTrue(settings.generate(MemoryOp.LONG_TERM_GET)); + assertTrue(settings.generate(MemoryOp.LONG_TERM_SEARCH)); + assertFalse(settings.generate(MemoryOp.SHORT_TERM_READ)); + assertFalse(settings.generate(MemoryOp.SENSORY_READ)); + } + + @Test + void testMasterSwitchFalseDisablesEverything() { + MemoryEventSettings settings = + MemoryEventSettings.from(Map.of("memory.generate-event", false)); + for (MemoryOp op : MemoryOp.values()) { + assertFalse(settings.generate(op)); + } + assertFalse(settings.anyEnabled()); + } + + @Test + void testMasterSwitchTrueEnablesReadsToo() { + MemoryEventSettings settings = + MemoryEventSettings.from(Map.of("memory.generate-event", true)); + assertTrue(settings.generate(MemoryOp.SHORT_TERM_READ)); + assertTrue(settings.generate(MemoryOp.SENSORY_READ)); + } + + @Test + void testSubKeyOverridesMaster() { + // YAML also accepts string-valued booleans. + MemoryEventSettings settings = + MemoryEventSettings.from( + Map.of( + "memory.generate-event", false, + "memory.generate-event.long-term-update", true, + "memory.generate-event.short-term-write", "true")); + assertTrue(settings.generate(MemoryOp.LONG_TERM_UPDATE)); + assertTrue(settings.generate(MemoryOp.SHORT_TERM_WRITE)); + // The explicit master value takes precedence over the per-operation default. + assertFalse(settings.generate(MemoryOp.SENSORY_WRITE)); + } + + @Test + void testExplicitFalseSubKeyBeatsMasterTrue() { + MemoryEventSettings settings = + MemoryEventSettings.from( + Map.of( + "memory.generate-event", true, + "memory.generate-event.long-term-search", false)); + assertFalse(settings.generate(MemoryOp.LONG_TERM_SEARCH)); + assertTrue(settings.generate(MemoryOp.LONG_TERM_GET)); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/TestMemoryObjectImplReadRecording.java b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/TestMemoryObjectImplReadRecording.java new file mode 100644 index 0000000000..c7ff0fd607 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/TestMemoryObjectImplReadRecording.java @@ -0,0 +1,124 @@ +/* + * 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.runtime.memory; + +import org.apache.flink.agents.api.context.MemoryObject; +import org.junit.jupiter.api.Test; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Tests read recording and {@link MemoryObjectImpl.MemoryItem#isValue()}. */ +public class TestMemoryObjectImplReadRecording { + + private MemoryObjectImpl newStm(@Nullable List readObservations) + throws Exception { + return new MemoryObjectImpl( + MemoryObject.MemoryType.SHORT_TERM, + new CachedMemoryStore(new ForTestMemoryMapState<>()), + MemoryObjectImpl.ROOT_KEY, + () -> {}, + new ArrayList<>(), + readObservations); + } + + @Test + void testGetValueRecordsAbsolutePathWhenEnabled() throws Exception { + List reads = new ArrayList<>(); + MemoryObjectImpl stm = newStm(reads); + stm.set("user.tier", "gold"); + + Object v = stm.get("user.tier").getValue(); + + assertEquals("gold", v); + assertEquals(1, reads.size()); + assertEquals("user.tier", reads.get(0).getPath()); + assertEquals("gold", reads.get(0).getValue()); + } + + @Test + void testGetValueWorksWithoutReadObservationSink() throws Exception { + MemoryObjectImpl stm = newStm(null); + stm.set("user.tier", "gold"); + + assertEquals("gold", stm.get("user.tier").getValue()); + } + + @Test + void testNestedChildInheritsReadRecording() throws Exception { + List reads = new ArrayList<>(); + MemoryObjectImpl stm = newStm(reads); + stm.set("user.address.city", "SF"); + stm.get("user").get("address").get("city").getValue(); + assertEquals(1, reads.size()); + assertEquals("user.address.city", reads.get(0).getPath()); + } + + @Test + void testGetFieldsRecordsOnlyValueChildren() throws Exception { + List reads = new ArrayList<>(); + MemoryObjectImpl stm = newStm(reads); + stm.newObject("empty", false); + stm.set("tier", "gold"); + + stm.getFields(); + + assertEquals(1, reads.size()); + assertEquals("tier", reads.get(0).getPath()); + assertEquals("gold", reads.get(0).getValue()); + } + + @Test + void testMemoryItemIsValue() { + // MemoryItem ctors are package-private; this test lives in the same package. + assertTrue(new MemoryObjectImpl.MemoryItem("x").isValue()); // VALUE item + assertFalse(new MemoryObjectImpl.MemoryItem().isValue()); // OBJECT item + } + + @Test + void testNestedChildInheritsMailboxChecker() throws Exception { + // Child objects share the parent's store and must also share its mailbox-thread contract. + AtomicInteger checkerCalls = new AtomicInteger(); + MemoryObjectImpl stm = + new MemoryObjectImpl( + MemoryObject.MemoryType.SHORT_TERM, + new CachedMemoryStore(new ForTestMemoryMapState<>()), + MemoryObjectImpl.ROOT_KEY, + checkerCalls::incrementAndGet, + new ArrayList<>(), + null); + stm.set("user.city", "SF"); + + MemoryObject child = stm.get("user"); // nested child object + int before = checkerCalls.get(); + child.getValue(); // op on the child must invoke the inherited checker (+1) + child.get("city"); // navigation on the child too (+1) + + assertEquals( + before + 2, + checkerCalls.get(), + "nested child must invoke the parent's real mailbox checker, not a no-op"); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java index c42dc02b12..2826a9bb45 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java @@ -27,9 +27,11 @@ import org.apache.flink.agents.api.context.DurableCallable; import org.apache.flink.agents.api.context.MemoryObject; import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.ShortTermWriteEvent; import org.apache.flink.agents.api.listener.EventListener; import org.apache.flink.agents.api.logger.EventLoggerConfig; import org.apache.flink.agents.api.logger.LoggerType; +import org.apache.flink.agents.api.memory.MemorySet; import org.apache.flink.agents.plan.AgentConfiguration; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.JavaFunction; @@ -39,6 +41,7 @@ import org.apache.flink.agents.runtime.actionstate.CallResult; import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore; import org.apache.flink.agents.runtime.eventlog.FileEventLogger; +import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory; import org.apache.flink.api.common.typeinfo.TypeInformation; import org.apache.flink.api.java.functions.KeySelector; import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; @@ -53,9 +56,12 @@ import org.junit.jupiter.api.Test; import java.io.IOException; +import java.io.Serializable; import java.lang.reflect.Field; +import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.UUID; @@ -67,6 +73,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.assertj.core.api.Assertions.catchThrowable; /** Tests for {@link ActionExecutionOperator}. */ public class ActionExecutionOperatorTest { @@ -77,6 +84,7 @@ public class ActionExecutionOperatorTest { void resetReconcilableFixtures() { TestAgent.resetReconcilableRecoveryFixture(); TestAgent.resetMixedRecoveryFixture(); + TestAgent.FOLLOWING_ACTION_EXECUTED.set(false); } @Test @@ -321,6 +329,46 @@ void pendingActionConfigSurvivesRestore() throws Exception { } } + @Test + void testRestoredActionUsesSameTextualContextKeyForLtmWriteAndCleanup() throws Exception { + AgentPlan agentPlan = TestAgent.getFailedActionAfterLtmAgentPlan(); + OperatorSubtaskState snapshot; + long key = 1L << 32; + String expectedContextKey = "4294967296"; + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(agentPlan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + testHarness.processElement(new StreamRecord<>(key)); + assertThat(testHarness.getTaskMailbox().size()).isEqualTo(1); + snapshot = testHarness.snapshot(1L, 1L); + } + + RecordingMem0LongTermMemory ltm = new RecordingMem0LongTermMemory(); + try (KeyedOneInputStreamOperatorTestHarness restoredHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(agentPlan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + restoredHarness.initializeState(snapshot); + restoredHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) restoredHarness.getOperator(); + replaceOperatorLtm(operator, ltm); + + assertThatThrownBy(operator::waitInFlightEventsFinished) + .hasCauseInstanceOf(ActionExecutionOperator.ActionTaskExecutionException.class) + .rootCause() + .hasMessageContaining("first action failed after LTM"); + + assertThat(ltm.recordedKeys()).containsExactly(expectedContextKey); + assertThat(ltm.drainedObservationKeys()).containsExactly(expectedContextKey); + } + } + @Test void testMemoryAccessProhibitedOutsideMailboxThread() throws Exception { try (KeyedOneInputStreamOperatorTestHarness testHarness = @@ -361,6 +409,216 @@ void testMailboxSubmittedActionTaskPropagatesError() throws Exception { } } + @Test + void testUnsupportedObservationValueIsSkippedWithoutChangingActionSuccess() throws Exception { + InMemoryActionStateStore actionStateStore = new InMemoryActionStateStore(false); + AgentPlan agentPlan = TestAgent.getBestEffortMemoryObservationPlan(); + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(agentPlan, true, actionStateStore), + (KeySelector) String::valueOf, + TypeInformation.of(String.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(7L)); + operator.waitInFlightEventsFinished(); + + List> output = + (List>) testHarness.getRecordOutput(); + assertThat(output).singleElement().extracting(StreamRecord::getValue).isEqualTo(7L); + + ActionState actionState = + actionStateStore.get( + "7", + 0L, + agentPlan.getActions().get("bestEffortMemoryObservationAction"), + new InputEvent(7L)); + assertThat(actionState).isNotNull(); + assertThat(actionState.getOutputEvents()) + .filteredOn(ShortTermWriteEvent.class::isInstance) + .singleElement() + .extracting(event -> ((ShortTermWriteEvent) event).getValue()) + .isEqualTo(Map.of("valid", 7)); + } + } + + @Test + void testFailedActionAfterLtmDiscardsCurrentKeyBeforeRethrowing() throws Exception { + RecordingMem0LongTermMemory ltm = new RecordingMem0LongTermMemory(); + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory( + TestAgent.getFailedActionAfterLtmAgentPlan(), true), + (KeySelector) String::valueOf, + TypeInformation.of(String.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + replaceOperatorLtm(operator, ltm); + + testHarness.processElement(new StreamRecord<>(0L)); + + assertThatThrownBy(() -> operator.waitInFlightEventsFinished()) + .hasCauseInstanceOf(ActionExecutionOperator.ActionTaskExecutionException.class) + .rootCause() + .hasMessageContaining("first action failed after LTM"); + // The public LTM operation ran before the action failed. + assertThat(ltm.recordedKeys()).containsExactly("0"); + // The common failure boundary explicitly drains this key before rethrowing. Check the + // same buffer directly rather than duplicating the Python LTM record schema here. + assertThat(ltm.pendingObservationKeys()).isEmpty(); + assertThat(ltm.drainedObservationKeys()).containsExactly("0"); + assertThat(ltm.drainCallCount()).isEqualTo(1); + // The failed mailbox task cannot continue to the following action in this operator. + assertThat(TestAgent.FOLLOWING_ACTION_EXECUTED).isFalse(); + } + } + + @Test + void testDiscardFailureDoesNotReplaceActionFailure() throws Exception { + RecordingMem0LongTermMemory ltm = new RecordingMem0LongTermMemory(); + RuntimeException discardFailure = new RuntimeException("discard failed"); + ltm.failDrainWith(discardFailure); + + try (KeyedOneInputStreamOperatorTestHarness testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory( + TestAgent.getFailedActionAfterLtmAgentPlan(), true), + (KeySelector) String::valueOf, + TypeInformation.of(String.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + replaceOperatorLtm(operator, ltm); + + testHarness.processElement(new StreamRecord<>(0L)); + Throwable thrown = catchThrowable(operator::waitInFlightEventsFinished); + + assertThat(thrown) + .hasCauseInstanceOf(ActionExecutionOperator.ActionTaskExecutionException.class) + .rootCause() + .hasMessageContaining("first action failed after LTM"); + assertThat(findSuppressedFailure(thrown, discardFailure)).isSameAs(discardFailure); + assertThat(ltm.recordedKeys()).containsExactly("0"); + assertThat(ltm.drainCallCount()).isEqualTo(1); + assertThat(TestAgent.FOLLOWING_ACTION_EXECUTED).isFalse(); + } + } + + private static Throwable findSuppressedFailure(Throwable failure, Throwable expected) { + for (Throwable current = failure; current != null; current = current.getCause()) { + for (Throwable suppressed : current.getSuppressed()) { + if (suppressed == expected) { + return suppressed; + } + } + } + return null; + } + + private static void replaceOperatorLtm( + ActionExecutionOperator operator, Mem0LongTermMemory ltm) throws Exception { + Field ltmField = ActionExecutionOperator.class.getDeclaredField("ltm"); + ltmField.setAccessible(true); + ltmField.set(operator, ltm); + } + + /** Java-side stand-in for the Python-backed LTM wrapper used to observe the failure path. */ + private static final class RecordingMem0LongTermMemory extends Mem0LongTermMemory { + private final List recordedKeys = new ArrayList<>(); + private final List pendingObservationKeys = new ArrayList<>(); + private final List pendingObservationIds = new ArrayList<>(); + private final List drainedObservationKeys = new ArrayList<>(); + private String currentKey; + private String currentObservationId; + private boolean updateObservationConfigured = true; + private boolean observationSuppressed; + private int drainCallCount; + private RuntimeException drainFailure; + + private RecordingMem0LongTermMemory() { + super(null, null); + } + + @Override + public void configureObservation( + boolean updateObservationEnabled, + boolean getObservationEnabled, + boolean searchObservationEnabled) { + updateObservationConfigured = updateObservationEnabled; + } + + @Override + public void switchContext( + String partitionKey, String observationId, boolean observationSuppressed) { + currentKey = partitionKey; + currentObservationId = observationId; + this.observationSuppressed = observationSuppressed; + } + + @Override + public MemorySet getMemorySet(String name) { + MemorySet memorySet = new MemorySet(name); + memorySet.setLtm(this); + return memorySet; + } + + @Override + public List add( + MemorySet memorySet, + List memoryItems, + @javax.annotation.Nullable List> metadatas) { + recordedKeys.add(currentKey); + if (!updateObservationConfigured || observationSuppressed) { + return List.of("memory-id"); + } + pendingObservationKeys.add(currentKey); + pendingObservationIds.add(currentObservationId); + return List.of("memory-id"); + } + + @Override + public String drainObservationRecordsJson(String partitionKey, String observationId) { + drainCallCount++; + if (drainFailure != null) { + throw drainFailure; + } + for (int index = pendingObservationKeys.size() - 1; index >= 0; index--) { + if (pendingObservationKeys.get(index).equals(partitionKey) + && pendingObservationIds.get(index).equals(observationId)) { + drainedObservationKeys.add(pendingObservationKeys.remove(index)); + pendingObservationIds.remove(index); + } + } + return "[]"; + } + + @Override + public void close() {} + + private List recordedKeys() { + return recordedKeys; + } + + private int drainCallCount() { + return drainCallCount; + } + + private List pendingObservationKeys() { + return pendingObservationKeys; + } + + private List drainedObservationKeys() { + return drainedObservationKeys; + } + + private void failDrainWith(RuntimeException failure) { + drainFailure = failure; + } + } + @Test void testInMemoryActionStateStoreIntegration() throws Exception { AgentPlan agentPlanWithStateStore = TestAgent.getAgentPlan(false); @@ -1608,6 +1866,9 @@ public static class TestAgent { public static final java.util.concurrent.atomic.AtomicInteger ACTION2_CALL_COUNTER = new java.util.concurrent.atomic.AtomicInteger(0); + public static final java.util.concurrent.atomic.AtomicBoolean FOLLOWING_ACTION_EXECUTED = + new java.util.concurrent.atomic.AtomicBoolean(false); + public static class MiddleEvent extends Event { public static final String EVENT_TYPE = "MiddleEvent"; @@ -1660,6 +1921,23 @@ public static void action3(MiddleEvent event, RunnerContext context) { } } + public static void bestEffortMemoryObservationAction(Event event, RunnerContext context) { + Long input = (Long) InputEvent.fromEvent(event).getInput(); + try { + context.getShortTermMemory().set("valid", input); + context.getShortTermMemory().set("kryo-only", new KryoOnlyObservationValue()); + context.getShortTermMemory().set("non-finite", Double.NaN); + context.sendEvent(new OutputEvent(input)); + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + } + + private static final class KryoOnlyObservationValue implements Serializable { + private static final long serialVersionUID = 1L; + private final Object value = new Object(); + } + private static DurableCallable durableCallable( String id, Class resultClass, Callable callSupplier) { return new DurableCallable() { @@ -1917,6 +2195,21 @@ private static boolean shouldThrowLinkageError() { public static void action(Event event, RunnerContext context) {} } + public static void failingActionAfterLtm(Event event, RunnerContext context) { + try { + context.getLongTermMemory() + .getMemorySet("notes") + .add(List.of("previous action"), null); + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + throw new IllegalStateException("first action failed after LTM"); + } + + public static void followingAction(Event event, RunnerContext context) { + FOLLOWING_ACTION_EXECUTED.set(true); + } + public static void resetReconcilableRecoveryFixture() { RECONCILABLE_CALL_COUNTER.set(0); RECONCILABLE_RECONCILE_COUNTER.set(0); @@ -2164,6 +2457,52 @@ public static AgentPlan getLinkageErrorAgentPlan() { return null; } + public static AgentPlan getBestEffortMemoryObservationPlan() { + try { + Action action = + new Action( + "bestEffortMemoryObservationAction", + new JavaFunction( + TestAgent.class, + "bestEffortMemoryObservationAction", + new Class[] {Event.class, RunnerContext.class}), + Collections.singletonList(InputEvent.EVENT_TYPE)); + return new AgentPlan(Map.of(action.getName(), action)); + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + return null; + } + + public static AgentPlan getFailedActionAfterLtmAgentPlan() { + try { + Map actions = new LinkedHashMap<>(); + Action failingAction = + new Action( + "failingActionAfterLtm", + new JavaFunction( + TestAgent.class, + "failingActionAfterLtm", + new Class[] {Event.class, RunnerContext.class}), + Collections.singletonList(InputEvent.EVENT_TYPE)); + Action followingAction = + new Action( + "followingAction", + new JavaFunction( + TestAgent.class, + "followingAction", + new Class[] {Event.class, RunnerContext.class}), + Collections.singletonList(InputEvent.EVENT_TYPE)); + actions.put(failingAction.getName(), failingAction); + actions.put(followingAction.getName(), followingAction); + + return new AgentPlan(actions); + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + return null; + } + // ==================== Actions for Exception Recovery Tests ==================== /** diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java index 8fec7cb63a..c5894df0f3 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java @@ -26,6 +26,7 @@ import org.apache.flink.agents.runtime.async.ContinuationContext; import org.apache.flink.agents.runtime.context.JavaRunnerContextImpl; import org.apache.flink.agents.runtime.context.RunnerContextImpl; +import org.apache.flink.agents.runtime.memory.InteranlBaseLongTermMemory; import org.apache.flink.agents.runtime.memory.MemoryObjectImpl; import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import org.apache.flink.api.common.state.MapState; @@ -134,6 +135,23 @@ void createAndSetRunnerContextReusesExistingMemoryContext() throws Exception { } } + @Test + void sameKeyTasksSwitchLtmWithDistinctObservationIds() throws Exception { + try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { + Action action = TestActions.noopAction(); + ActionTask suspended = new JavaActionTask("k", new InputEvent(1L), action); + ActionTask sibling = new JavaActionTask("k", new InputEvent(1L), action); + InteranlBaseLongTermMemory ltm = mock(InteranlBaseLongTermMemory.class); + + invokeCreateAndSetRunnerContext(mgr, suspended, ltm); + invokeCreateAndSetRunnerContext(mgr, sibling, ltm); + + assertThat(suspended.getObservationId()).isNotEqualTo(sibling.getObservationId()); + verify(ltm).switchContext("k", suspended.getObservationId(), false); + verify(ltm).switchContext("k", sibling.getObservationId(), false); + } + } + @Test void transferContextsCopiesMemoryAndContinuationToNewTask() throws Exception { try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { @@ -229,6 +247,14 @@ void closeIsIdempotent() throws Exception { @SuppressWarnings("unchecked") private static void invokeCreateAndSetRunnerContext( ActionTaskContextManager mgr, ActionTask task) { + invokeCreateAndSetRunnerContext(mgr, task, null); + } + + @SuppressWarnings("unchecked") + private static void invokeCreateAndSetRunnerContext( + ActionTaskContextManager mgr, + ActionTask task, + InteranlBaseLongTermMemory longTermMemory) { AgentPlan plan = newEmptyAgentPlan(); ResourceCache cache = mock(ResourceCache.class); FlinkAgentsMetricGroupImpl metricGroup = @@ -246,7 +272,7 @@ private static void invokeCreateAndSetRunnerContext( sensoryMem, shortTermMem, /* pythonRunnerContext */ null, - /* longTermMemory */ null); + longTermMemory); } private static AgentPlan newEmptyAgentPlan() { diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java index 998998ba6d..90d451a9c5 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java @@ -19,17 +19,98 @@ import org.apache.flink.agents.api.Event; import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.plan.PythonFunction; import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl; +import org.apache.flink.agents.runtime.python.operator.PythonActionTask; +import org.apache.flink.agents.runtime.python.operator.PythonGeneratorActionTask; +import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor; import org.junit.jupiter.api.Test; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; /** Tests for output finalization owned by {@link ActionTask}. */ class ActionTaskTest { + @Test + void distinctActionExecutionsHaveDistinctObservationIds() { + Event triggeringEvent = new InputEvent(1L); + Action action = TestActions.noopAction(); + + ActionTask first = new JavaActionTask("key", triggeringEvent, action); + ActionTask second = new JavaActionTask("key", triggeringEvent, action); + + assertThat(first.getObservationId()).isNotEqualTo(second.getObservationId()); + assertThat(first).isNotEqualTo(second); + } + + @Test + void taskRestoredWithoutObservationIdGetsStableFallback() { + ActionTask task = new JavaActionTask("key", new InputEvent(1L), TestActions.noopAction()); + task.observationId = null; + + String restoredObservationId = task.getObservationId(); + + assertThat(restoredObservationId).isNotBlank(); + assertThat(task.getObservationId()).isEqualTo(restoredObservationId); + } + + @Test + void pythonContinuationKeepsObservationId() throws Exception { + Action action = + new Action( + "python-action", + new PythonFunction("test_module", "test_action"), + List.of(InputEvent.EVENT_TYPE)); + PythonActionTask task = new PythonActionTask("key", new InputEvent(1L), action); + PythonRunnerContextImpl context = mock(PythonRunnerContextImpl.class); + PythonActionExecutor executor = mock(PythonActionExecutor.class); + task.setRunnerContext(context); + when(context.getPythonAwaitableRef()).thenReturn("awaitable"); + when(executor.executePythonFunction((PythonFunction) action.getExec(), task.event)) + .thenReturn("awaitable"); + when(executor.callPythonAwaitable("awaitable")).thenReturn(false); + + ActionTask generated = + task.invoke(getClass().getClassLoader(), executor) + .getGeneratedActionTask() + .orElseThrow(); + + assertThat(generated).isInstanceOf(PythonGeneratorActionTask.class); + assertThat(generated.getObservationId()).isEqualTo(task.getObservationId()); + } + + @Test + void pythonRecoveredContinuationKeepsObservationId() throws Exception { + Action action = + new Action( + "python-action", + new PythonFunction("test_module", "test_action"), + List.of(InputEvent.EVENT_TYPE)); + PythonGeneratorActionTask task = + new PythonGeneratorActionTask("key", new InputEvent(1L), action, "observation-id"); + PythonRunnerContextImpl context = mock(PythonRunnerContextImpl.class); + PythonActionExecutor executor = mock(PythonActionExecutor.class); + task.setRunnerContext(context); + when(context.getPythonAwaitableRef()).thenReturn(null, "restarted-awaitable"); + when(executor.executePythonFunction((PythonFunction) action.getExec(), task.event)) + .thenReturn("restarted-awaitable"); + when(executor.callPythonAwaitable("restarted-awaitable")).thenReturn(false); + + ActionTask generated = + task.invoke(getClass().getClassLoader(), executor) + .getGeneratedActionTask() + .orElseThrow(); + + assertThat(generated).isInstanceOf(PythonGeneratorActionTask.class); + assertThat(generated.getObservationId()).isEqualTo("observation-id"); + } + @Test void resultFinalizesOutputLineage() { Event triggeringEvent = new InputEvent(1L); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/AgentRunBeginEventEmissionTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/AgentRunBeginEventEmissionTest.java new file mode 100644 index 0000000000..528df47a9c --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/AgentRunBeginEventEmissionTest.java @@ -0,0 +1,303 @@ +/* + * 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.runtime.operator; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.EventContext; +import org.apache.flink.agents.api.EventType; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.agents.AgentExecutionOptions; +import org.apache.flink.agents.api.configuration.AgentConfigOptions; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.event.AgentRunBeginEvent; +import org.apache.flink.agents.api.listener.EventListener; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.plan.JavaFunction; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.apache.flink.util.ExceptionUtils; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for AgentRunBeginEvent emission by {@link ActionExecutionOperator}. */ +public class AgentRunBeginEventEmissionTest { + + /** Collects every processed event through the EVENT_LISTENERS seam. */ + public static class CollectingEventListener implements EventListener { + public static final List EVENTS = new CopyOnWriteArrayList<>(); + + @Override + public void onEventProcessed(EventContext context, Event event) { + EVENTS.add(event); + } + } + + private static final AtomicBoolean RUN_BEGIN_ACTION_EXECUTED = new AtomicBoolean(false); + + @BeforeEach + void reset() { + CollectingEventListener.EVENTS.clear(); + RUN_BEGIN_ACTION_EXECUTED.set(false); + } + + /** Input-triggered action: writes one STM value. */ + public static void writeTierAction(Event event, RunnerContext context) { + try { + context.getShortTermMemory().set("user.tier", "gold"); + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + } + + /** Run-begin-triggered action: writes an STM flag and records execution. */ + public static void onRunBeginAction(Event event, RunnerContext context) { + try { + context.getShortTermMemory().set("run.begin.seen", true); + context.sendEvent(new OutputEvent("run-begin")); + RUN_BEGIN_ACTION_EXECUTED.set(true); + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + } + + private static AgentPlan buildPlan(AgentConfiguration config, boolean subscribeRunBegin) { + try { + Map actions = new HashMap<>(); + Action inputAction = + new Action( + "writeTierAction", + new JavaFunction( + AgentRunBeginEventEmissionTest.class, + "writeTierAction", + new Class[] {Event.class, RunnerContext.class}), + Collections.singletonList(InputEvent.EVENT_TYPE)); + actions.put(inputAction.getName(), inputAction); + + if (subscribeRunBegin) { + Action runBeginAction = + new Action( + "onRunBeginAction", + new JavaFunction( + AgentRunBeginEventEmissionTest.class, + "onRunBeginAction", + new Class[] {Event.class, RunnerContext.class}), + Collections.singletonList(EventType.AgentRunBeginEvent)); + actions.put(runBeginAction.getName(), runBeginAction); + } + return new AgentPlan(actions, new HashMap<>(), config); + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + return null; + } + + private static AgentConfiguration listenerConfig() { + AgentConfiguration config = new AgentConfiguration(); + config.set( + AgentConfigOptions.EVENT_LISTENERS, + List.of(CollectingEventListener.class.getName())); + return config; + } + + private static AgentConfiguration listenerConfigWithRunBeginEvent() { + AgentConfiguration config = listenerConfig(); + config.set(AgentExecutionOptions.AGENT_RUN_BEGIN_EVENT, true); + return config; + } + + private static KeyedOneInputStreamOperatorTestHarness createHarness( + AgentPlan agentPlan) throws Exception { + return new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(agentPlan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class)); + } + + private static List collectedRunBeginEvents() { + return CollectingEventListener.EVENTS.stream() + .filter(e -> AgentRunBeginEvent.EVENT_TYPE.equals(e.getType())) + .map(AgentRunBeginEvent::fromEvent) + .collect(Collectors.toList()); + } + + @Test + void testRunBeginEmittedOncePerInputWithPreviousRunStm() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness testHarness = + createHarness(buildPlan(listenerConfigWithRunBeginEvent(), false))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(42L)); + operator.waitInFlightEventsFinished(); + + List runBegins = collectedRunBeginEvents(); + assertThat(runBegins).hasSize(1); + assertThat(runBegins.get(0).getKey()).isEqualTo("42"); + assertThat(runBegins.get(0).getValue()).isEmpty(); + assertThat(runBegins.get(0).getId()).isNotNull(); + + testHarness.processElement(new StreamRecord<>(42L)); + operator.waitInFlightEventsFinished(); + + runBegins = collectedRunBeginEvents(); + assertThat(runBegins).hasSize(2); + assertThat(runBegins.get(1).getKey()).isEqualTo("42"); + assertThat(runBegins.get(1).getValue()).containsEntry("user.tier", "gold"); + assertThat(runBegins.get(1).getId()).isNotEqualTo(runBegins.get(0).getId()); + } + } + + @Test + void testRunBeginCopiesInputSourceTimestamp() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness testHarness = + createHarness(buildPlan(listenerConfigWithRunBeginEvent(), true))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(42L, 1234L)); + operator.waitInFlightEventsFinished(); + + Event input = + CollectingEventListener.EVENTS.stream() + .filter(event -> InputEvent.EVENT_TYPE.equals(event.getType())) + .findFirst() + .orElseThrow(); + assertThat(collectedRunBeginEvents()) + .singleElement() + .satisfies( + event -> { + assertThat(event.hasSourceTimestamp()).isTrue(); + assertThat(event.getSourceTimestamp()).isEqualTo(1234L); + assertThat(event.getUpstreamEventId()).isEqualTo(input.getId()); + assertThat(event.getUpstreamActionName()) + .isEqualTo("agent_run_begin_action"); + }); + List> output = + (List>) testHarness.getRecordOutput(); + assertThat(output) + .singleElement() + .satisfies( + record -> { + assertThat(record.hasTimestamp()).isTrue(); + assertThat(record.getTimestamp()).isEqualTo(1234L); + }); + } + } + + @Test + void testRunBeginSnapshotIsolatedPerKey() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness testHarness = + createHarness(buildPlan(listenerConfigWithRunBeginEvent(), false))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + // key 1 writes user.tier; key 2's own run-begin must not see key 1's STM. + testHarness.processElement(new StreamRecord<>(1L)); + operator.waitInFlightEventsFinished(); + testHarness.processElement(new StreamRecord<>(2L)); + operator.waitInFlightEventsFinished(); + testHarness.processElement(new StreamRecord<>(1L)); + operator.waitInFlightEventsFinished(); + + List runBegins = collectedRunBeginEvents(); + assertThat(runBegins).hasSize(3); + // key 2's first run-begin sees only its own (empty) STM, no leak from key 1. + assertThat(runBegins.get(1).getKey()).isEqualTo("2"); + assertThat(runBegins.get(1).getValue()).isEmpty(); + // key 1's second run-begin carries key 1's own prior write. + assertThat(runBegins.get(2).getKey()).isEqualTo("1"); + assertThat(runBegins.get(2).getValue()).containsEntry("user.tier", "gold"); + } + } + + @Test + void testRunBeginDisabledByDefault() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness testHarness = + createHarness(buildPlan(listenerConfig(), false))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(42L)); + operator.waitInFlightEventsFinished(); + + assertThat(collectedRunBeginEvents()).isEmpty(); + } + } + + @Test + void testRunBeginDisabledByConfig() throws Exception { + AgentConfiguration config = listenerConfig(); + config.set(AgentExecutionOptions.AGENT_RUN_BEGIN_EVENT, false); + try (KeyedOneInputStreamOperatorTestHarness testHarness = + createHarness(buildPlan(config, false))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(42L)); + operator.waitInFlightEventsFinished(); + + assertThat(collectedRunBeginEvents()).isEmpty(); + } + } + + @Test + void testRunBeginEventTriggersSubscribedAction() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness testHarness = + createHarness(buildPlan(listenerConfigWithRunBeginEvent(), true))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + + testHarness.processElement(new StreamRecord<>(7L)); + operator.waitInFlightEventsFinished(); + + // Proves routing through getActionsTriggeredBy, not just logging. + assertThat(RUN_BEGIN_ACTION_EXECUTED).isTrue(); + + // The subscribed action's STM write shows up in the NEXT run's begin snapshot. + testHarness.processElement(new StreamRecord<>(7L)); + operator.waitInFlightEventsFinished(); + + List runBegins = collectedRunBeginEvents(); + assertThat(runBegins).hasSize(2); + assertThat(runBegins.get(1).getValue()) + .containsEntry("run.begin.seen", true) + .containsEntry("user.tier", "gold"); + } + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java index 72e26fa44d..f4af793315 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/EventLineageReconstructionTest.java @@ -22,6 +22,7 @@ import org.apache.flink.agents.api.InputEvent; import org.apache.flink.agents.api.OutputEvent; import org.apache.flink.agents.api.configuration.AgentConfigOptions; +import org.apache.flink.agents.api.configuration.MemoryEventOptions; import org.apache.flink.agents.api.logger.EventLogLevel; import org.apache.flink.agents.api.logger.LoggerType; import org.apache.flink.agents.plan.AgentConfiguration; @@ -104,6 +105,7 @@ private static AgentConfiguration fileLoggerConfig(Path logDir) { config.set(AgentConfigOptions.BASE_LOG_DIR, logDir.toString()); config.set(AgentConfigOptions.EVENT_LOG_LEVEL, EventLogLevel.STANDARD); config.set(AgentConfigOptions.EVENT_LOG_MAX_STRING_LENGTH, 3); + config.set(MemoryEventOptions.MEMORY_GENERATE_EVENT, false); return config; } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java new file mode 100644 index 0000000000..baf584fa4f --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java @@ -0,0 +1,94 @@ +/* + * 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.runtime.python.utils; + +import org.apache.flink.types.Row; +import org.junit.jupiter.api.Test; +import pemja.core.PythonInterpreter; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class PythonActionExecutorTest { + + @Test + void resolvesPickledPythonKeyTextFromPyFlinkKeyRow() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonActionExecutor executor = newExecutor(interpreter); + byte[] pickledKey = new byte[] {1, 2, 3}; + when(interpreter.invoke( + "python_java_utils.convert_to_python_key_text", pickledKey, "pickled")) + .thenReturn("7"); + + assertThat(executor.resolveKeyText(Row.of(pickledKey), true)).isEqualTo("7"); + verify(interpreter) + .invoke("python_java_utils.convert_to_python_key_text", pickledKey, "pickled"); + } + + @Test + void resolvesExplicitPyFlinkKeyTypesWithStringValueOf() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonActionExecutor executor = newExecutor(interpreter); + + assertThat(executor.resolveKeyText(Row.of(7L), false)).isEqualTo("7"); + assertThat(executor.resolveKeyText(Row.of(42), false)).isEqualTo("42"); + } + + @Test + void resolvesExplicitByteArrayWithoutUnpickling() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonActionExecutor executor = newExecutor(interpreter); + byte[] firstKey = new byte[] {'N', '.'}; + byte[] secondKey = new byte[] {(byte) 0x80, 0x04, 'N', '.'}; + when(interpreter.invoke( + "python_java_utils.convert_to_python_key_text", firstKey, "explicit")) + .thenReturn("b'N.'"); + when(interpreter.invoke( + "python_java_utils.convert_to_python_key_text", secondKey, "explicit")) + .thenReturn("b'\\x80\\x04N.'"); + + assertThat(executor.resolveKeyText(Row.of(firstKey), false)).isEqualTo("b'N.'"); + assertThat(executor.resolveKeyText(Row.of(secondKey), false)).isEqualTo("b'\\x80\\x04N.'"); + verify(interpreter) + .invoke("python_java_utils.convert_to_python_key_text", firstKey, "explicit"); + verify(interpreter) + .invoke("python_java_utils.convert_to_python_key_text", secondKey, "explicit"); + } + + @Test + void propagatesDecodeFailure() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonActionExecutor executor = newExecutor(interpreter); + byte[] malformedKey = new byte[] {2}; + when(interpreter.invoke( + "python_java_utils.convert_to_python_key_text", malformedKey, "pickled")) + .thenThrow(new RuntimeException("bad pickle")); + + assertThatThrownBy(() -> executor.resolveKeyText(Row.of(malformedKey), true)) + .isInstanceOf(RuntimeException.class) + .hasMessage("bad pickle"); + } + + private static PythonActionExecutor newExecutor(PythonInterpreter interpreter) + throws Exception { + return new PythonActionExecutor(interpreter, null, null, null, "test-job"); + } +}