[api][runtime] Support Short-Term Memory (#67)

Author: KeGu-069

api/src/main/java/org/apache/flink/agents/api/context/MemoryObject.java

@@ -0,0 +1,104 @@
+/*
+ * 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.context;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * A representation of an object in the short-term memory. It is responsible for accessing and
+ * manipulating (direct or indirect) fields within the memory structure. A direct field is a field
+ * which stores primitive data directly, while an indirect filed is a field which represents a
+ * nested object.Fields can be accessed using an absolute or relative path.
+ */
+public interface MemoryObject {
+    /**
+     * Returns a MemoryObject that represents the given path.
+     *
+     * @param path relative path from the current object to the target field
+     * @return a MemoryObject instance pointing to the field. If the field is a primitive value
+     *     type, the value of the returned MemoryObject can be exposed via {@link #getValue()}.If
+     *     the field is a nested object, the subfields of the returned MemoryObject can be exposed
+     *     via {@link #getFields()}.
+     * @throws Exception if the field does not exist
+     */
+    MemoryObject get(String path) throws Exception;
+
+    /**
+     * Sets the value of a direct field in the current object; any missing intermediate objects will
+     * be created automatically.
+     *
+     * @param path relative path from the current object to the target field
+     * @param value new value of the field
+     * @throws Exception if trying to overwrite a nested object with a primitive value or set a
+     *     MemoryObject directly
+     */
+    void set(String path, Object value) throws Exception;
+
+    /**
+     * Creates a new MemoryObject as an indirect field in the current object.
+     *
+     * @param path relative path from the current object to the target field
+     * @param overwrite whether to overwrite existing field if it's not a nested object
+     * @return the created object
+     * @throws Exception if field exists but is not a nested object and overwrite is false
+     */
+    MemoryObject newObject(String path, boolean overwrite) throws Exception;
+
+    /**
+     * Checks whether a (direct or indirect) field exists in the current object.
+     *
+     * @param path relative path from the current object to the target field
+     * @return true if the field exists, false otherwise
+     */
+    boolean isExist(String path);
+
+    /**
+     * Gets names of all the top-level fields of the current object.
+     *
+     * @return list of top-level field names
+     * @throws Exception state-backend failure
+     */
+    List<String> getFieldNames() throws Exception;
+
+    /**
+     * Gets all the top-level fields of the current object.
+     *
+     * @return map of top-level fields
+     * @throws Exception state-backend failure
+     */
+    Map<String, Object> getFields() throws Exception;
+
+    /**
+     * Gets the primitive value stored at the current path.
+     *
+     * @return the primitive value if this MemoryObject stores primitive value directly, or null if
+     *     it represents a nested object
+     * @throws Exception state-backend failure
+     */
+    Object getValue() throws Exception;
+
+    /**
+     * Checks whether the current object is a nested object.
+     *
+     * @return true if this MemoryObject is a nested object, false if it stores primitive value
+     *     directly
+     * @throws Exception state-backend failure
+     */
+    boolean isNestedObject() throws Exception;
+}

api/src/main/java/org/apache/flink/agents/api/context/RunnerContext.java

@@ -30,4 +30,12 @@ public interface RunnerContext {
      * @param event the event to be sent
      */
     void sendEvent(Event event);
+
+    /**
+     * Gets the short-term memory.
+     *
+     * @return MemoryObject the root of the short-term memory
+     * @throws Exception if the underlying state backend cannot be accessed
+     */
+    MemoryObject getShortTermMemory() throws Exception;
 }

python/flink_agents/api/memory_object.py

@@ -0,0 +1,110 @@
+################################################################################
+#  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.
+#################################################################################
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+from pydantic import BaseModel
+
+
+class MemoryObject(BaseModel, ABC):
+    """Representation of an object in the short-term memory.
+
+    A direct field is a field which stores concrete data directly, while an indirect
+    filed is just a field which represents a nested object. Fields can be accessed
+    using an absolute or relative path.
+    """
+
+    @abstractmethod
+    def get(self, path: str) -> Any:
+        """Get the value of a (direct or indirect) field in the object.
+
+        Parameters
+        ----------
+        path: str
+          Relative path from the current object to the target field.
+
+        Returns:
+        -------
+        Any
+          If the field is a direct field, returns the concrete data stored.
+          If the field is an indirect field, another MemoryObject will be returned.
+          If the field doesn't exist, returns None.
+        """
+
+    @abstractmethod
+    def set(self, path: str, value: Any) -> None:
+        """Set the value of an indirect field in the object.
+        This will also create the intermediate objects if not exist.
+
+        Parameters
+        ----------
+        path: str
+          Relative path from the current object to the target field.
+        value: Any
+          New value of the field. The type of the value must be a primary type.
+        """
+
+    @abstractmethod
+    def new_object(self, path: str) -> "MemoryObject":
+        """Create a new object as the value of an indirect field in the object.
+
+        Parameters
+        ----------
+        path: str
+          Relative path from the current object to the target field.
+
+        Returns:
+        -------
+        MemoryObject
+          The created object.
+        """
+
+    @abstractmethod
+    def is_exist(self, path: str) -> bool:
+        """Check whether a (direct or indirect) field exist in the object.
+
+        Parameters
+        ----------
+        path: str
+          Relative path from the current object to the target field.
+
+        Returns:
+        -------
+        bool
+          Whether the field exists.
+        """
+
+    @abstractmethod
+    def get_field_names(self) -> List[str]:
+        """Get names of all the top-level subfields of the object.
+
+        Returns:
+        -------
+        List[str]
+          Top-level subfield names of the object in a list.
+        """
+
+    @abstractmethod
+    def get_fields(self) -> Dict[str, Any]:
+        """Get all the top-level subfields of the object.
+
+        Returns:
+        -------
+        Dict[str, Any]
+          Top-level subfields of the object in a dictionary.
+        """

python/flink_agents/api/runner_context.py

@@ -18,6 +18,7 @@
 from abc import ABC, abstractmethod
 
 from flink_agents.api.event import Event
+from flink_agents.api.memory_object import MemoryObject
 from flink_agents.api.resource import Resource, ResourceType
 
 
@@ -48,3 +49,13 @@ def get_resource(self, name: str, type: ResourceType) -> Resource:
         type : ResourceType
             The type of the resource.
         """
+
+    @abstractmethod
+    def get_short_term_memory(self) -> MemoryObject:
+        """Get the short-term memory.
+
+        Returns:
+        -------
+        MemoryObject
+          The root object of the short-term memory.
+        """

python/flink_agents/examples/agent_example.py

@@ -40,16 +40,32 @@ class MyAgent(Agent):
     @staticmethod
     def first_action(event: Event, ctx: RunnerContext):  # noqa D102
         input = event.input
-        content = input + " first_action"
+        memory = ctx.get_short_term_memory()
+
+        current_count = memory.get("action_counter") or 0
+        new_count = current_count + 1
+        memory.set("action_counter", new_count)
+
+        content = input + " -> first_action"
+        key_with_count = f"(seen {new_count} times)"
+
         ctx.send_event(MyEvent(value=content))
-        ctx.send_event(OutputEvent(output=content))
+        ctx.send_event(OutputEvent(output={key_with_count: content}))
 
     @action(MyEvent)
     @staticmethod
     def second_action(event: Event, ctx: RunnerContext):  # noqa D102
         input = event.value
-        content = input + " second_action"
-        ctx.send_event(OutputEvent(output=content))
+        memory = ctx.get_short_term_memory()
+
+        current_count = memory.get("action_counter")
+        new_count = current_count + 1
+        memory.set("action_counter", new_count)
+
+        base_message = input.split("->")[0].strip()
+        content = base_message + " -> second_action"
+        key_with_count = f"(seen {new_count} times)"
+        ctx.send_event(OutputEvent(output={key_with_count: content}))
 
 
 if __name__ == "__main__":
@@ -62,8 +78,10 @@ def second_action(event: Event, ctx: RunnerContext):  # noqa D102
 
     input_list.append({"key": "bob", "value": "The message from bob"})
     input_list.append({"k": "john", "v": "The message from john"})
+    input_list.append({"key": "john", "value": "Second message from john"})
+    input_list.append({"key": "bob", "value": "Second message from bob"})
     input_list.append(
-        {"value": "The message from unknown"}
+        {"value": "Message from unknown"}
     )  # will automatically generate a new unique key
 
     env.execute()

python/flink_agents/examples/my_agent.py

@@ -16,7 +16,7 @@
 # limitations under the License.
 #################################################################################
 import copy
-from typing import Any
+from typing import Any, Optional
 
 from pydantic import BaseModel
 
@@ -42,6 +42,7 @@ class ItemData(BaseModel):
     id: int
     review: str
     review_score: float
+    memory_info: Optional[dict] = None
 
 
 class MyEvent(Event):  # noqa D101
@@ -60,6 +61,16 @@ class DataStreamAgent(Agent):
     @staticmethod
     def first_action(event: Event, ctx: RunnerContext):  # noqa D102
         input = event.input
+
+        stm = ctx.get_short_term_memory()
+        status = stm.new_object("status", overwrite = True)
+
+        total = 0
+        if stm.is_exist("status.total_reviews"):
+            total = status.get("total_reviews")
+        total += 1
+        status.set("total_reviews", total)
+
         content = copy.deepcopy(input)
         content.review += " first action"
         ctx.send_event(MyEvent(value=content))
@@ -68,8 +79,15 @@ def first_action(event: Event, ctx: RunnerContext):  # noqa D102
     @staticmethod
     def second_action(event: Event, ctx: RunnerContext):  # noqa D102
         input = event.value
+
+        stm = ctx.get_short_term_memory()
+        memory_info = {
+            "total_reviews": stm.get("status.total_reviews"),
+        }
+
         content = copy.deepcopy(input)
         content.review += " second action"
+        content.memory_info = memory_info
         ctx.send_event(OutputEvent(output=content))