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 f486a6d0e..661673cea 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 @@ -40,5 +40,23 @@ 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 EventType() {} } 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 a0b9269e4..74e3cf967 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 000000000..5cc31b013 --- /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 000000000..1c886436a --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/AgentRunBeginEvent.java @@ -0,0 +1,74 @@ +/* + * 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 String key and before any action of that run executes. Carries the key'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) { + AgentRunBeginEvent result = new AgentRunBeginEvent(event.getId(), event.getAttributes()); + if (event.hasSourceTimestamp()) { + result.setSourceTimestamp(event.getSourceTimestamp()); + } + return result; + } + + @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 000000000..7c6d09fa6 --- /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 000000000..3d5ddca69 --- /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 000000000..00f1d556f --- /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 000000000..24f29bc46 --- /dev/null +++ b/api/src/main/java/org/apache/flink/agents/api/event/MemoryEvent.java @@ -0,0 +1,138 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.api.event; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.core.json.JsonWriteFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import org.apache.flink.agents.api.Event; + +import java.io.IOException; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.UUID; +import java.util.function.BiFunction; + +/** + * Base class of the memory observation events. One event per (memory scope x operation) is emitted + * at the action finish boundary; each concrete subclass pins one of the seven operation kinds as + * its {@code type}. + * + *

Attributes: {@code key} — the String Flink key the operation belongs to; {@code value} — the + * operation's folded JSON value map. Framework observation events are skipped for non-String keyed + * streams. 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()); + } + MemoryEvent result = factory.apply(event.getId(), event.getAttributes()); + if (event.hasSourceTimestamp()) { + result.setSourceTimestamp(event.getSourceTimestamp()); + } + return result; + } + + @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 000000000..a077aa07d --- /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 000000000..6b420d86d --- /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 000000000..8c2be4cfa --- /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 000000000..bc7b79a8f --- /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 5ae8bc7e8..b09e6a13a 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/event/AgentRunBeginEventTest.java b/api/src/test/java/org/apache/flink/agents/api/event/AgentRunBeginEventTest.java new file mode 100644 index 000000000..57bc75a4b --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/event/AgentRunBeginEventTest.java @@ -0,0 +1,131 @@ +/* + * 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 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 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 000000000..72c9e7959 --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/event/MemoryEventTest.java @@ -0,0 +1,259 @@ +/* + * 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 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 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/memory_events.md b/docs/content/docs/development/memory/memory_events.md new file mode 100644 index 000000000..c7cd5ccf2 --- /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. 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 only for String-keyed streams and carry two common attributes: + +| Attribute | Description | +|-----------|-------------| +| `key` | The String Flink key whose memory was observed. Non-String keys are skipped safely. | +| `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.** The default `ON_READ_AND_WRITE` policy refreshes TTL on reads. If you enable `agent-run.begin-event`, its snapshot scan refreshes TTL for visited entries. Use `ON_CREATE_AND_WRITE` when entries should expire based on writes only. +{{< /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 a3b8b5e8c..2ab5db675 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 >}} +The default `ON_READ_AND_WRITE` update type extends an entry's lifetime whenever it is read. This also applies when producing the run-begin snapshot used by [Memory Events]({{< ref "docs/development/memory/memory_events" >}}): if you opt in through `agent-run.begin-event`, each input scans the key's short-term memory and refreshes TTL for the entries it reads, although only value nodes are included in the event. Choose `ON_CREATE_AND_WRITE` when entries should expire based only on writes. +{{< /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 44ebbc0fe..78fd9495d 100644 --- a/docs/content/docs/operations/configuration.md +++ b/docs/content/docs/operations/configuration.md @@ -147,9 +147,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 000000000..5c050f0ac --- /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 000000000..a6f906c2f --- /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 000000000..79626d1d4 --- /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 000000000..1a4ac90a9 --- /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 d196777c4..7acb082ae 100644 --- a/python/flink_agents/api/core_options.py +++ b/python/flink_agents/api/core_options.py @@ -215,6 +215,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.""" @@ -248,6 +281,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 344d6167c..d1635b032 100644 --- a/python/flink_agents/api/events/event_type.py +++ b/python/flink_agents/api/events/event_type.py @@ -37,6 +37,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, ) @@ -59,3 +83,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 000000000..18d1fbebf --- /dev/null +++ b/python/flink_agents/api/events/memory_event.py @@ -0,0 +1,205 @@ +################################################################################ +# 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) + result.id = event.id + return result + + @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 000000000..dbc6ddc5e --- /dev/null +++ b/python/flink_agents/api/events/run_event.py @@ -0,0 +1,62 @@ +################################################################################ +# 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") + ) + result.id = event.id + return result + + @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 315301aeb..d147b233d 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 @@ -165,7 +167,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 @@ -191,8 +195,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 ─────────────────────────────────────────────────── @@ -223,7 +227,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." @@ -237,7 +243,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) @@ -340,7 +350,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" @@ -377,12 +389,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" @@ -391,6 +407,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 000000000..eb9a5b14f --- /dev/null +++ b/python/flink_agents/api/tests/test_memory_event.py @@ -0,0 +1,234 @@ +################################################################################ +# 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_before_id_generation() -> None: + original = {"nested": {"value": 1}, "bytes": b"\x01\x02\x03"} + event = ShortTermWriteEvent(key="k", value=original) + same_wire = ShortTermWriteEvent( + key="k", value={"nested": {"value": 1}, "bytes": "AQID"} + ) + + original["nested"]["value"] = 2 + assert event.value == {"nested": {"value": 1}, "bytes": "AQID"} + assert event.id == same_wire.id + + +@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 000000000..1dadeaeea --- /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/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 000000000..2bed19ffa --- /dev/null +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/memory_event_logging_test.py @@ -0,0 +1,117 @@ +################################################################################ +# 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: str(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 "key" in write["event"]["attributes"] + + begins = [r for r in records if r["eventType"] == EventType.AgentRunBeginEvent] + assert len(begins) == len(inputs) + # 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 0e2762d73..a1750a3f6 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/memory/mem0/mem0_long_term_memory.py b/python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py index bcd81dfd3..39e8ac975 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,13 @@ class Mem0LongTermMemory(InternalBaseLongTermMemory): _mem0: Any = PrivateAttr(default=None) + _ltm_observation_records: queue.Queue[tuple[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) + def __init__( self, *, @@ -224,7 +270,13 @@ def _create_mem0_instance(self) -> Any: return Memory(mem0_config) @override - def switch_context(self, key: str) -> None: + def switch_context( + self, + key: str, + update_observation_enabled: bool | None = None, + get_observation_enabled: bool | None = None, + search_observation_enabled: bool | None = None, + ) -> None: """Switch the keyed partition context. This method is called on the mailbox thread before each action. @@ -233,12 +285,126 @@ def switch_context(self, key: str) -> None: Args: key: The new key for partition isolation. + update_observation_enabled: Whether mutations should be recorded. + get_observation_enabled: Whether gets should be recorded. + search_observation_enabled: Whether searches should be recorded. """ # Ensure Mem0 is initialized on the mailbox thread. _ = self._mem0_instance # Ensure report token usage on the mailbox thread self._report_token_metrics() self.key = key + if update_observation_enabled is not None: + self._update_observation_enabled = update_observation_enabled + if get_observation_enabled is not None: + self._get_observation_enabled = get_observation_enabled + if search_observation_enabled is not None: + self._search_observation_enabled = search_observation_enabled + + def _record_ltm_op( + self, + op: _LtmObservationOp | str, + memory_set: str, + mem_id: str | None, + value: Any, + observation_key: str, + *, + enabled: bool = True, + ) -> None: + """Buffer one LTM operation for observation. + + Args: + op: The operation kind. + memory_set: The name of the memory set operated on. + mem_id: The affected memory id, or None for whole-set ops. + value: The stored memory content, or None when not applicable. + observation_key: Partition key captured at operation entry. + enabled: Whether this operation type is configured for observation. + """ + try: + if not enabled or getattr(self.ctx, "_j_runner_context", None) is None: + return + try: + operation = _LtmObservationOp(op) + except ValueError: + logger.warning("Skipping unknown LTM observation operation %r", op) + return + self._ltm_observation_records.put( + ( + observation_key, + _LtmObservationRecord( + op=operation.value, + set=memory_set, + id=mem_id, + value=value, + ), + ) + ) + except Exception: + logger.debug("LTM observation buffering failed; skipping", exc_info=True) + + def _record_ltm_search( + self, + memory_set: str, + query: str, + hits: List[Dict[str, Any]], + observation_key: str, + *, + enabled: bool = True, + ) -> None: + """Buffer one LTM search call for observation. + + Args: + memory_set: The set searched. + query: The search query string. + hits: The ordered matched records, each with id, value, and score. + observation_key: Partition key captured at operation entry. + enabled: Whether search observation is configured. + """ + try: + if not enabled or getattr(self.ctx, "_j_runner_context", None) is None: + return + self._ltm_observation_records.put( + ( + observation_key, + _LtmObservationRecord( + op=_LtmObservationOp.SEARCH.value, + set=memory_set, + query=query, + value=hits, + ), + ) + ) + except Exception: + logger.debug("LTM observation buffering failed; skipping", exc_info=True) + + def drain_ltm_observation_records(self, key: str) -> str: + """Pop buffered LTM records for one partition key as a JSON array. + + Called from Java on the mailbox thread at action-finish flush. Records for + other partition keys are placed back in the shared queue. + + Args: + key: Partition key whose records to drain. + + Returns: + JSON array string of the drained records. + """ + records: List[_LtmObservationRecord] = [] + other_records: List[tuple[str, _LtmObservationRecord]] = [] + while True: + try: + owner_key, record = self._ltm_observation_records.get_nowait() + except queue.Empty: + break + if owner_key == key: + records.append(record) + else: + other_records.append((owner_key, 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 +428,21 @@ def delete_memory_set(self, name: str) -> bool: Returns: True if the memory set was deleted. """ + observation_key = self.key + observation_enabled = self._update_observation_enabled self._mem0_instance.delete_all( user_id=self.job_id, agent_id=self.key, run_id=name, ) + self._record_ltm_op( + _LtmObservationOp.DELETE_SET, + name, + None, + None, + observation_key, + enabled=observation_enabled, + ) return True @override @@ -286,6 +462,8 @@ def add( Returns: List of IDs of the added memories. """ + observation_key = self.key + observation_enabled = self._update_observation_enabled if isinstance(memory_items, str): memory_items = [memory_items] if metadatas is not None and isinstance(metadatas, dict): @@ -302,9 +480,27 @@ def add( 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, + enabled=observation_enabled, + ) return all_ids @override @@ -327,6 +523,8 @@ def get( Returns: List of memory items. """ + observation_key = self.key + observation_enabled = self._get_observation_enabled if ids is not None: if isinstance(ids, str): ids = [ids] @@ -334,6 +532,15 @@ 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, + enabled=observation_enabled, + ) return items result = self._mem0_instance.get_all( @@ -343,10 +550,20 @@ def get( 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, + enabled=observation_enabled, + ) + return items @override def delete(self, memory_set: MemorySet, ids: str | List[str] | None = None) -> None: @@ -356,18 +573,36 @@ 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_enabled = self._update_observation_enabled if ids is None: self._mem0_instance.delete_all( user_id=self.job_id, agent_id=self.key, run_id=memory_set.name, ) + self._record_ltm_op( + _LtmObservationOp.DELETE_SET, + memory_set.name, + None, + None, + observation_key, + 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, + enabled=observation_enabled, + ) @override def search( @@ -390,6 +625,8 @@ def search( Returns: List of matching memory items. """ + observation_key = self.key + observation_enabled = self._search_observation_enabled result = self._mem0_instance.search( query=query, user_id=self.job_id, @@ -399,9 +636,20 @@ def search( 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, + 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 +703,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_op_recording.py b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py new file mode 100644 index 000000000..e1f865eda --- /dev/null +++ b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py @@ -0,0 +1,111 @@ +################################################################################ +# 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 typing import Any +from unittest.mock import MagicMock + +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 + ctx._j_runner_context = object() + ltm = Mem0LongTermMemory.model_construct( + ctx=ctx, job_id="job", key="partition", metric_group=None + ) + ltm._mem0 = mem0 + ltm._update_observation_enabled = True + ltm._get_observation_enabled = True + ltm._search_observation_enabled = True + return ltm + + +def _drain(ltm: Mem0LongTermMemory, key: str = "partition") -> list[dict]: + return json.loads(ltm.drain_ltm_observation_records(key)) + + +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_partition_key_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" + 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") == [] + + +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" 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 000000000..27ef899ec --- /dev/null +++ b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_recording_hook.py @@ -0,0 +1,119 @@ +################################################################################ +# 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 + ctx._j_runner_context = object() + 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") + ltm._record_ltm_search("policies", "refund", [{"id": "p1"}], "partition") + + assert json.loads(ltm.drain_ltm_observation_records("partition")) == [ + { + "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_rebuffers_other_partition_records() -> None: + ltm = _make_ltm() + ltm._record_ltm_op(_LtmObservationOp.ADD, "s", "a", "va", "partition-a") + ltm._record_ltm_op(_LtmObservationOp.ADD, "s", "b", "vb", "partition-b") + + assert [ + record["id"] + for record in json.loads(ltm.drain_ltm_observation_records("partition-a")) + ] == ["a"] + assert [ + record["id"] + for record in json.loads(ltm.drain_ltm_observation_records("partition-b")) + ] == ["b"] + + +def test_disabled_operation_does_not_enqueue_or_json_encode() -> None: + ltm = _make_ltm() + ltm._record_ltm_op( + _LtmObservationOp.ADD, + "s", + "m1", + object(), + "partition", + 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" + ) + + 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"))) == 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") 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 1395bdc56..472c38217 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,21 @@ public CachedMemoryStore getSensoryMemStore() { protected String actionName; protected InteranlBaseLongTermMemory ltm; + /** Supported String Flink key for emitted framework observation events. */ + @Nullable protected String eventKeyText; + + /** True when the current action was triggered by a memory event: suppress observation. */ + protected boolean observationSuppressed; + + /** Hashed partition key used to route LTM observation records. */ + protected String ltmPartitionKey; + + /** 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; + /** Context for fine-grained durable execution, may be null if not enabled. */ @Nullable protected DurableExecutionContext durableExecutionContext; @@ -118,17 +154,44 @@ public RunnerContextImpl( this.mailboxThreadChecker = mailboxThreadChecker; this.agentPlan = agentPlan; this.resourceCache = resourceCache; + this.memoryEventSettings = MemoryEventSettings.from(agentPlan.getConfigData()); } public void setLongTermMemory(InteranlBaseLongTermMemory ltm) { this.ltm = ltm; } - public void switchActionContext(String actionName, MemoryContext memoryContext, String key) { + public void switchActionContext( + String actionName, + MemoryContext memoryContext, + String ltmPartitionKey, + @Nullable String eventKeyText, + boolean observationSuppressed) { this.actionName = actionName; this.memoryContext = memoryContext; + this.ltmPartitionKey = ltmPartitionKey; + this.eventKeyText = eventKeyText; + this.observationSuppressed = observationSuppressed; + boolean observationAllowed = !observationSuppressed && eventKeyText != null; + boolean updateObservationEnabled = + observationAllowed + && memoryEventSettings.generate( + MemoryEventSettings.MemoryOp.LONG_TERM_UPDATE); + boolean getObservationEnabled = + observationAllowed + && memoryEventSettings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_GET); + boolean searchObservationEnabled = + observationAllowed + && memoryEventSettings.generate( + MemoryEventSettings.MemoryOp.LONG_TERM_SEARCH); + this.ltmObservationEnabled = + updateObservationEnabled || getObservationEnabled || searchObservationEnabled; if (ltm != null) { - ltm.switchContext(key); + ltm.switchContext( + ltmPartitionKey, + updateObservationEnabled, + getObservationEnabled, + searchObservationEnabled); } } @@ -161,6 +224,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 || ltmPartitionKey == null || !ltmObservationEnabled) { + return; + } + ltm.drainObservationRecordsJson(ltmPartitionKey); + } + + private List drainPendingEvents(Long timestamp) { List list = new ArrayList<>(this.pendingEvents); if (timestamp != null) { list.forEach(event -> event.setSourceTimestamp(timestamp)); @@ -169,6 +258,60 @@ 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; + } + if (eventKeyText == null) { + LOG.warn( + "Skipping framework memory observation for action '{}' because the Flink key is not a String", + actionName); + return; + } + + List> ltmRecords = Collections.emptyList(); + if (ltm != null && ltmObservationEnabled) { + try { + ltmRecords = + MemoryEventBuilder.parseLtmObservationRecords( + ltm.drainObservationRecordsJson(ltmPartitionKey)); + } catch (Exception | LinkageError e) { + LOG.warn( + "LTM observation drain failed for action '{}' and partition key '{}' ({}); skipping records", + actionName, + ltmPartitionKey, + e.getClass().getSimpleName()); + } + } + try { + pendingEvents.addAll( + MemoryEventBuilder.buildWriteEvents( + eventKeyText, + memoryContext.getSensoryMemoryUpdates(), + memoryContext.getShortTermMemoryUpdates(), + memoryEventSettings)); + pendingEvents.addAll( + MemoryEventBuilder.buildReadEvents( + eventKeyText, sensoryReads, shortTermReads, memoryEventSettings)); + pendingEvents.addAll( + MemoryEventBuilder.buildLtmEvents( + eventKeyText, 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 +337,37 @@ public List getShortTermMemoryUpdates() { @Override public MemoryObject getSensoryMemory() throws Exception { mailboxThreadChecker.run(); + List memoryReads = null; + if (eventKeyText != null + && !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 (eventKeyText != null + && !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 b13e2d7ca..396477591 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 @@ -25,7 +25,20 @@ public interface InteranlBaseLongTermMemory extends BaseLongTermMemory { * 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 */ - void switchContext(String key); + void switchContext( + String partitionKey, + boolean updateObservationEnabled, + boolean getObservationEnabled, + boolean searchObservationEnabled); + + /** + * Drains the LTM observation records buffered for one partition key, 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 + * @return JSON array string of drained records + */ + String drainObservationRecordsJson(String partitionKey); } 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 2847812a4..bd17082e1 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,30 @@ public List search( } @Override - public void switchContext(String key) { - adapter.callMethod(pyMem0, "switch_context", Map.of("key", key)); + public void switchContext( + String partitionKey, + boolean updateObservationEnabled, + boolean getObservationEnabled, + boolean searchObservationEnabled) { + adapter.callMethod( + pyMem0, + "switch_context", + Map.of( + "key", + partitionKey, + "update_observation_enabled", + updateObservationEnabled, + "get_observation_enabled", + getObservationEnabled, + "search_observation_enabled", + searchObservationEnabled)); + } + + @Override + public String drainObservationRecordsJson(String partitionKey) { + return (String) + adapter.callMethod( + pyMem0, "drain_ltm_observation_records", Map.of("key", partitionKey)); } @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 000000000..a1839c680 --- /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 000000000..e83a15e8a --- /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 b9550a676..80d49c548 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 000000000..4e987d298 --- /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 46d09e0a2..0dac746e0 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; @@ -118,6 +126,9 @@ 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); } @@ -263,6 +278,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, event); } // We then obtain the triggered action and add ActionTasks to the waiting processing // queue. @@ -280,6 +296,54 @@ private void processEvent(Object key, Event event) throws Exception { } } + /** + * Attempts to emit an {@link AgentRunBeginEvent} for the input before any action triggered by + * that input executes. + */ + private void tryEmitAgentRunBeginEvent(Object key, Event inputEvent) throws Exception { + if (!agentRunBeginEventEnabled) { + return; + } + String eventKeyText = resolveEventKeyText(key); + if (eventKeyText == null) { + LOG.warn( + "Skipping AgentRunBeginEvent because framework observation requires a String Flink key"); + 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(eventKeyText, 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()); + } + processEvent(key, beginEvent); + } + private void tryProcessActionTaskForKey(Object key) { try { processActionTaskForKey(key); @@ -319,6 +383,7 @@ private void processActionTaskForKey(Object key) throws Exception { contextManager.createAndSetRunnerContext( actionTask, key, + resolveEventKeyText(key), agentPlan, resourceCache, metricGroup, @@ -373,10 +438,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 @@ -555,6 +633,27 @@ private ActionTask createActionTask(Object key, Action action, Event event) { } } + /** Returns the logical String key for Java and PyFlink keyed streams, if supported. */ + @Nullable + private String resolveEventKeyText(Object key) { + PythonActionExecutor pythonActionExecutor = + pythonBridge == null ? null : pythonBridge.getPythonActionExecutor(); + return resolveEventKeyText(key, inputIsJava, pythonActionExecutor); + } + + @VisibleForTesting + @Nullable + static String resolveEventKeyText( + Object key, boolean inputIsJava, @Nullable PythonActionExecutor pythonActionExecutor) { + if (key instanceof String) { + return (String) key; + } + if (inputIsJava || pythonActionExecutor == null) { + return null; + } + return pythonActionExecutor.resolveStringKey(key); + } + private void tryResumeProcessActionTasks() throws Exception { Iterable keys = stateManager.getProcessingKeys(); if (keys != null) { 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 eba25d5d4..4344609a5 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; @@ -157,6 +158,8 @@ RunnerContextImpl createOrGetRunnerContext( * * @param actionTask the task to be set up before execution. * @param key the current Flink key. + * @param eventKeyText the logical String key used by framework observation events, or {@code + * null} for unsupported key types. * @param agentPlan the agent plan. * @param resourceCache the resource cache. * @param metricGroup the agent metric group. @@ -170,6 +173,7 @@ RunnerContextImpl createOrGetRunnerContext( void createAndSetRunnerContext( ActionTask actionTask, Object key, + @Nullable String eventKeyText, AgentPlan agentPlan, ResourceCache resourceCache, FlinkAgentsMetricGroupImpl metricGroup, @@ -218,7 +222,11 @@ void createAndSetRunnerContext( } context.switchActionContext( - actionTask.action.getName(), memoryContext, String.valueOf(key.hashCode())); + actionTask.action.getName(), + memoryContext, + String.valueOf(key.hashCode()), + eventKeyText, + 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 11724ce68..19bd677f9 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 cea9cce54..219a07e6a 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 @@ -235,6 +235,8 @@ private void wireLongTermMemory() { LongTermMemoryOptions.Mem0.VECTOR_STORE.getKey())); } longTermMemory = new Mem0LongTermMemory(pythonResourceAdapter, (PyObject) pyLtm); + // The Python runner context drains LTM observation records at action finish. + pythonRunnerContext.setLongTermMemory(longTermMemory); } 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 ce43195a5..aeebbfeeb 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 @@ -63,6 +63,6 @@ public ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec 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 35296c1cc..fc42c83c9 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 @@ -56,7 +56,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 67c80f38d..e73a64250 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,9 +24,12 @@ 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; +import javax.annotation.Nullable; + import java.io.IOException; import java.util.concurrent.atomic.AtomicLong; @@ -64,6 +67,8 @@ 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_OBJECT = + "python_java_utils.convert_to_python_object"; 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"; @@ -161,6 +166,34 @@ public Event wrapToInputEvent(Object eventData) throws IOException { return Event.fromJson((String) result); } + /** Resolves a logical String key from PyFlink's keyed-stream representation. */ + @Nullable + public String resolveStringKey(Object flinkKey) { + Object logicalKey = flinkKey; + if (flinkKey instanceof Row) { + Row row = (Row) flinkKey; + if (row.getArity() != 1) { + return null; + } + logicalKey = row.getField(0); + } + + if (logicalKey instanceof String) { + return (String) logicalKey; + } + if (!(logicalKey instanceof byte[])) { + return null; + } + + try { + Object decoded = interpreter.invoke(CONVERT_TO_PYTHON_OBJECT, logicalKey); + return decoded instanceof String ? (String) decoded : null; + } catch (RuntimeException e) { + // Framework observation is best-effort. Unsupported or malformed keys are skipped. + return null; + } + } + 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/actionstate/ActionStateSerdeTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java index f8dbc6fa8..c35e1988b 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 000000000..769584039 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestMemoryObservationFlush.java @@ -0,0 +1,325 @@ +/* + * 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.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, @Nullable String eventKeyText) + throws Exception { + return createContext(conf, suppressed, eventKeyText, null); + } + + private RunnerContextImpl createContext( + Map conf, + boolean suppressed, + @Nullable String eventKeyText, + @Nullable InteranlBaseLongTermMemory ltm) + throws Exception { + AgentPlan plan = + new AgentPlan( + new HashMap<>(), + 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) { + context.setLongTermMemory(ltm); + } + context.switchActionContext( + "action", + new RunnerContextImpl.MemoryContext( + new CachedMemoryStore(new ForTestMemoryMapState<>()), + new CachedMemoryStore(new ForTestMemoryMapState<>())), + "legacy-partition-hash", + eventKeyText, + 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 suppressedOrNonStringKeyDoesNotCrossLtmBridge() throws Exception { + StubLtm suppressedLtm = new StubLtm(); + RunnerContextImpl suppressed = + createContext(new HashMap<>(), true, "user-42", suppressedLtm); + suppressed.drainEventsAtActionFinish(null); + suppressed.discardMemoryObservation(); + + StubLtm nonStringLtm = new StubLtm(); + RunnerContextImpl nonStringKey = createContext(new HashMap<>(), false, null, nonStringLtm); + nonStringKey.drainEventsAtActionFinish(null); + nonStringKey.discardMemoryObservation(); + + assertThat(suppressedLtm.drainCallCount).isZero(); + assertThat(nonStringLtm.drainCallCount).isZero(); + } + + @Test + void nonStringKeyDoesNotRecordMemoryObservationsButStillPersistsUpdates() 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, null); + + context.getSensoryMemory().set("sensory", "value"); + context.getSensoryMemory().get("sensory").getValue(); + context.getShortTermMemory().set("short-term", "value"); + context.getShortTermMemory().get("short-term").getValue(); + + assertReadObservationListsEmpty(context.getMemoryContext()); + assertThat(context.getSensoryMemoryUpdates()).hasSize(1); + assertThat(context.getShortTermMemoryUpdates()).hasSize(1); + assertThat(context.drainEventsAtActionFinish(null)).isEmpty(); + } + + @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 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 int drainCallCount; + private LinkageError drainLinkageError; + + @Override + public void switchContext( + String partitionKey, + boolean updateObservationEnabled, + boolean getObservationEnabled, + boolean searchObservationEnabled) { + this.updateEnabled = updateObservationEnabled; + this.getEnabled = getObservationEnabled; + this.searchEnabled = searchObservationEnabled; + } + + @Override + public String drainObservationRecordsJson(String partitionKey) { + 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() {} + } +} 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 916d68cd6..1abec5d37 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 eee167bdf..f00dbbf6c 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,38 @@ 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.switchContext("k1", true, false, true); + ltm.drainObservationRecordsJson("k1"); ltm.close(); verify(mockAdapter) - .callMethod(eq(mockPyMem0), eq("switch_context"), eq(Map.of("key", "k1"))); + .callMethod( + eq(mockPyMem0), + eq("switch_context"), + eq( + Map.of( + "key", + "k1", + "update_observation_enabled", + true, + "get_observation_enabled", + false, + "search_observation_enabled", + true))); + verify(mockAdapter) + .callMethod( + eq(mockPyMem0), + eq("drain_ltm_observation_records"), + eq(Map.of("key", "k1"))); 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 97e4099e2..54a42e0c8 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 000000000..1139d5c4d --- /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 000000000..48acb524e --- /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 000000000..c7ff0fd60 --- /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 22c0072f8..d0d945cce 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 @@ -26,9 +26,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; @@ -37,6 +39,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; @@ -50,7 +53,9 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.io.Serializable; import java.lang.reflect.Field; +import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -63,6 +68,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 { @@ -73,6 +79,7 @@ public class ActionExecutionOperatorTest { void resetReconcilableFixtures() { TestAgent.resetReconcilableRecoveryFixture(); TestAgent.resetMixedRecoveryFixture(); + TestAgent.FOLLOWING_ACTION_EXECUTED.set(false); } @Test @@ -310,6 +317,203 @@ 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(String.valueOf("0".hashCode())); + // 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(String.valueOf("0".hashCode())); + 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(String.valueOf("0".hashCode())); + 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 drainedObservationKeys = new ArrayList<>(); + private String currentKey; + private boolean updateObservationEnabled; + private int drainCallCount; + private RuntimeException drainFailure; + + private RecordingMem0LongTermMemory() { + super(null, null); + } + + @Override + public void switchContext( + String partitionKey, + boolean updateObservationEnabled, + boolean getObservationEnabled, + boolean searchObservationEnabled) { + currentKey = partitionKey; + this.updateObservationEnabled = updateObservationEnabled; + } + + @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 (!updateObservationEnabled) { + return List.of("memory-id"); + } + pendingObservationKeys.add(currentKey); + return List.of("memory-id"); + } + + @Override + public String drainObservationRecordsJson(String partitionKey) { + drainCallCount++; + if (drainFailure != null) { + throw drainFailure; + } + while (pendingObservationKeys.remove(partitionKey)) { + drainedObservationKeys.add(partitionKey); + } + 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); @@ -1439,6 +1643,9 @@ public static class TestAgent { public static final java.util.concurrent.atomic.AtomicInteger DURABLE_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"; @@ -1489,6 +1696,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() { @@ -1746,6 +1970,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); @@ -1995,6 +2234,57 @@ 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), + Map.of(InputEvent.EVENT_TYPE, List.of(action)), + new HashMap<>()); + } catch (Exception e) { + ExceptionUtils.rethrow(e); + } + return null; + } + + public static AgentPlan getFailedActionAfterLtmAgentPlan() { + try { + Map> actionsByEvent = new HashMap<>(); + Map actions = new HashMap<>(); + 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)); + actionsByEvent.put(InputEvent.EVENT_TYPE, List.of(failingAction, followingAction)); + actions.put(failingAction.getName(), failingAction); + actions.put(followingAction.getName(), followingAction); + + return new AgentPlan(actions, actionsByEvent, new HashMap<>()); + } 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 8fec7cb63..a7efe36a5 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 @@ -238,6 +238,7 @@ private static void invokeCreateAndSetRunnerContext( mgr.createAndSetRunnerContext( task, "k", + "k", plan, cache, metricGroup, 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 000000000..2dcb4fb6d --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/AgentRunBeginEventEmissionTest.java @@ -0,0 +1,314 @@ +/* + * 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.agents.runtime.python.utils.PythonActionExecutor; +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.types.Row; +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; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verifyNoInteractions; + +/** 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> actionsByEvent = new HashMap<>(); + 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)); + actionsByEvent.put(InputEvent.EVENT_TYPE, Collections.singletonList(inputAction)); + 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)); + actionsByEvent.put( + EventType.AgentRunBeginEvent, Collections.singletonList(runBeginAction)); + actions.put(runBeginAction.getName(), runBeginAction); + } + return new AgentPlan(actions, actionsByEvent, 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) String::valueOf, + TypeInformation.of(String.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(); + + assertThat(collectedRunBeginEvents()) + .singleElement() + .satisfies( + event -> { + assertThat(event.hasSourceTimestamp()).isTrue(); + assertThat(event.getSourceTimestamp()).isEqualTo(1234L); + }); + 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"); + } + } + + @Test + void javaRowKeyIsNotDecodedAsPyFlinkKey() { + PythonActionExecutor pythonActionExecutor = mock(PythonActionExecutor.class); + + assertThat( + ActionExecutionOperator.resolveEventKeyText( + Row.of(new byte[] {1, 2, 3}), true, pythonActionExecutor)) + .isNull(); + verifyNoInteractions(pythonActionExecutor); + } +} 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 000000000..498787410 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java @@ -0,0 +1,64 @@ +/* + * 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.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class PythonActionExecutorTest { + + @Test + void resolvesPickledStringFromPyFlinkKeyRow() 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_object", pickledKey)) + .thenReturn("user-7"); + + assertThat(executor.resolveStringKey(Row.of(pickledKey))).isEqualTo("user-7"); + verify(interpreter).invoke("python_java_utils.convert_to_python_object", pickledKey); + } + + @Test + void rejectsNonStringAndMalformedPyFlinkKeys() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonActionExecutor executor = newExecutor(interpreter); + byte[] nonStringKey = new byte[] {1}; + byte[] malformedKey = new byte[] {2}; + when(interpreter.invoke("python_java_utils.convert_to_python_object", nonStringKey)) + .thenReturn(7); + when(interpreter.invoke("python_java_utils.convert_to_python_object", malformedKey)) + .thenThrow(new RuntimeException("bad pickle")); + + assertThat(executor.resolveStringKey(Row.of(nonStringKey))).isNull(); + assertThat(executor.resolveStringKey(Row.of(malformedKey))).isNull(); + assertThat(executor.resolveStringKey(Row.of("too", "many"))).isNull(); + assertThat(executor.resolveStringKey(7)).isNull(); + } + + private static PythonActionExecutor newExecutor(PythonInterpreter interpreter) + throws Exception { + return new PythonActionExecutor(interpreter, null, null, null, "test-job"); + } +}