Short-Term Memory Design

[xintongsong]

Clarification of Concepts

• Run: An agent run refers to the processing of a single input event.
• Key: In Flink Agents, inputs of the agent are partitioned by their keys. This corresponds to how data are partitioned by keys in Flink's Keyed DataStream.
• Memory: Data stored during execution of agents, for remembering things from previous agent runs and recall them in subsequent agent runs.
  • Short-Term Memory: Memory shared across all actions within an agent run, and multiple agent runs with the same input key. This corresponds to Flink's Keyed State, which is visible to processing of multiple records with in the same keyed partition, and is transparent to processing of data in other keyed partitions.
  • Long-Term Memory: Memory shared across all keyed partitions. (Design of long-term memory is not in the scope of this documentation.)

Design Overview

The above figure shows the design of Short-Term Memory.

• Logically, the short-term memory may have multiple fields, where each field is either a primary type or an objection. For an object, it may have more fields of primary types or objects. The latter results in a chain of nested objects.
• Physically, we store the short-term memory in a Flink MapState. This requires flatten the nested objects into a single-layer map. We use the absolute path, the sequence of field names starting from the root (Short-Term Memory) that points to the current field, as the map key. For primary types, the map value is the field value. For objects, the map value is a special mark indicating the existence of the object. The mark also contains a list of direct field names of the object, so that when fetching the whole object, we don't need to iterate over the entire map state looking for keys with certain prefix.

Please notice that the root node in the figure is also an object. Therefore, the key to the short-term memory design is around storing and accessing of objects.

APIs

MemoryObject

We introduced MemoryObject as the interface for accessing an object in the short-term memory.

class MemoryObject(ABC):
  """
  Representation of an object in the short-term memory.                
  """

  @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
      The value of the field. If the field is an object, another MemoryObject will be returned. If the field doesn't exist, returns None.
    """

  @abstractmethod
  def set(self, path: str, value: Any):
    """
    Set the value of a (direct or 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 either a primary type, or MemoryObject.
    """

  @abstractmethod
  def new_object(self, path: str) -> MemoryObject:
    """
    Create a new object as 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:
    -------
    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 direct fields of the object.

    Returns:
    -------
    List[str]
      Direct field names of the object in a list.
    """

  @abstractmethod
  def get_fields(self) -> Dict[str, Any]:
    """
    Get all the direct fields of the object.

    Returns:
    -------
    Dict[str, Any]
      Direct fields in a dictionary.
    """

RunnerContext

class RunnerContext(ABC):
  # ...

  @abstractmethod
  def get_short_term_memory(self) -> MemoryObject:
    """
    Get the short-term memory.

    Returns:
    -------
    MemoryObject
      The root object of the short-term memory.
    """

Examples

@action(MyEvent)
@staticmethod
def my_action(event: Event, ctx: RunnerContext):
  # ...

  root = ctx.get_short_term_memory()

  root.set("x", 100)
  root.set("y", "abc")
  ojb_z = root.new_object("z")

  # {
  #   "x" : 100,
  #   "y" : "abc"
  #   "z" : {}
  # }

  obj_z.set("m", 0.5)
  obj_z.set("n.j", True)

  # {
  #   "x" : 100,
  #   "y" : "abc"
  #   "z" : {
  #     "m" : 0.5,
  #     "n" : {
  #       "j" : True
  #     }    
  #   }
  # }

  root.get("x") # 100
  root.get("z.m") # 0.5
  
  root.get("xx") # None
  root.get("z.mm") # None

  root.is_exist("x") # True
  root.is_exist("xx") # False
  root.is_exist("z.m") # True
  root.is_exist("z.mm") # False

  obj_n = root.get("z.n")
  obj_n.get("j") # True

  obj_z.get_field_names() # ["m", "n"]
  obj_z.get_fields() # {"m" : 0.5, "n" : MemoryObject}

[coderplay]

A couple of questions:

1. What is typically used as the partition key here? Could you give an example? For instance, would it be a host name in logs, or a user ID in a clickstream?

2. It looks like the data types for short-term memory are primitives and nested objects, correct? Do we support lists, sets , queues?

3. If we want to call back the history of the last 5 runs, how would we do that?

4. Flink state is persisted permanently and uploaded to blob storage after a checkpoint. What is the retention period for our short-term memory? How does the user define this retention?

5. Lastly but most importantly, conversation history is the most common form of short-term memory, but long conversations (in our case, which means large number of runs) are challenging for current LLMs. A full history may not fit within an LLM’s context window, leading to unrecoverable errors. Even if the LLM supports the full context length, most models still perform poorly over long contexts, becoming “distracted” by stale or off-topic content while also suffering from slower response times and higher costs. What mechanisms does our design provide to address this issue?

[xintongsong]

1. The key can be user / device / product id, similar to Flink's keyed data stream.

2. I think lists should be supported. For sets and queues, not sure about the necessity. Maybe we can start with lists, and add sets and queues later if they are proven needed?

3 & 4. Good question. Flink supports TTL of states, which expires old states automatically based on processing / event time. But I think what the agent need is something like keeping the most recent N messages, or keeping up to M tokens / bytes of history messages. Of course users can manually control the retention and expiration of states. But we'd better provide a easy way for doing that.

5. I think to address this issue, we would need some mechanism for conclude a long history into a brief summary, or retrieve most relevant pieces from a large history. However, I'm not entirely sure whether this should be done with Flink's short-term memory. Maybe leveraging an external vector storage or AI memory project would work better? Such information summary and relevant context retrieving usually involves understanding the semantics of history, embedding, gpu usage, etc., which are not what Flink's state backend is good at.

[coderplay]

A few more questions:

1. This design uses whether the memory is keyed as the criterion to distinguish short-term vs. long-term memory, but this is confusing. Even keyed state in Flink is permanently persisted.

2. Short-term memory is mainly used to synthesize context for LLM models, and this context is typically in natural language format. Thus, the most common memory types are str, list[str]. Supporting deeply nested object memory structures feels like over-engineering, right?

3. Why does memory have to be stored as MapState? When Flink stores MapState in StateBackends, it uses a three-level hierarchy:

• The input stream key, which defines task parallelism.
• The MapStateDescriptor name, which defines the map state's name.
• The key within the MapState.
  Using the StateDescriptor name alone should be sufficient to distinguish different memory entries, correct? Additionally, ListStateDescriptor can directly support list-type state.

[xintongsong]

1. This is a good point. I think we should not define a memory is short-term or long-term based on whether it is shared across keyed partitions. Visible to other keyed partitions or not is more like a property or restriction of the memory, rather than definition.

Offline discussed with @coderplay about this. Here are something we came up with.

• Short-Term Memory
  • Stores states and information into Flink's state backend
  • Accessed directly in a get-X style, e.g., get(field_name), get_recent(list_name, n)
  • Needs retention control supports for fields of collection types. When the collection size (number of elements, or total size in bytes) reaches a threshold, the framework should either remove the earliest / least-recently-used element or call a user defined aggregation function to reduce the collection size.
• Long-Term Memory
  • Stores in a vector storage
  • Accessed in a searching style
  • The vector storage can be either external or built-in (like rocksdb).
  • The vector storage should also support snapshots, to be integrated with Flink's checkpointing mechanism.
• We can also support automatically moving short-term memory into long-term memory, as an option of the retention control.

[xintongsong]

2. Short-term memory is used not only for providing context for the LLM, but also to share information and task status across actions within an agent. Nested structure is useful when it come to reusable actions.

Imagine an action that generates a short summary from a long text, and all inputs / outputs need to be stored in memory. In this case, it would be more efficient to have the action read input from and write output to the state rather than copying the text into events and sending between event.

The question is, if the action is called multiple times in an agent, which part of the state should it ready / write? We can define a memory object with two fields: original_text and summary. With the nested structure, every time we call the action, we can pass in the path of a different memory object, and the action can read the original_text field of the object from states, and write the summary field.

[xintongsong]

Why not using state descriptor to distinguish memory entries.

• Having multiple state descriptors is less efficient compared to multiple entries in a map state, in terms of per-state (descriptor) overhead, state / checkpoint management complexity, compression ratio, etc.
• The flink community is planning to drop the supports for dynamically creating new states in runtime, because it leads to many problems in state management and compatibility control. Instead, we will require users to declare all the states to be used before executing the job. That might be not flexible enough for short-term memory, compared to a map state where you can add new k-v entries anytime you want.

Why not using the list state

• A map state can cover most use cases that other types of state (value, list, ...) can cover. It is more flexible to support both list-types and other types or named status & properties.
• The map state comes with an advantage that when accessing an entry of the map we don't need to deserialize the whole state.
• AFAIK, the only advantage of the list states is that, it has better supports for redistribution (split and merge), which is helpful when using as an operator state and when we change the parallelism of a Flink job. Which is not relevant to us because we are using keyed states.

[KeGu-069]

Introduce

In Flink-Agents, actions pass relevant messages through events and share state through short-term memory. When large data (such as conversation history or complex custom object) need to be passed between multiple actions, passing data directly through events can incur significant performance overhead due to repeated serialization and deserialization.

To solve this problem, we introduce MemoryRef—a persistent reference to data in Short-Term Memory. It has two main benefits:

1. Reduced Transfer Overhead: When data is passed along an Action chain (e.g., A->B->C), but an intermediate Action (like B) does not need to use the data, passing a MemoryRef avoids the serialization/deserialization overhead of the actual data.
2. Data Persistence: A MemoryRef internally records the data's original type information. This allows us to store data of any type into the state and restore it to its original type when needed, thus enabling reliable persistence for arbitrary data types.
   By using MemoryRef, we can achieve efficient and low-cost transfer of any large or frequently referenced data between Actions.

Overview

A MemoryRef is a serializable data structure that does not store the actual memory data itself. Instead, it acts as a "pointer" or "handle" to a specific data item in Short-Term Memory.

A MemoryRef instance contains just two key pieces of information:

• path: The absolute path to the referenced data in Short-Term Memory (a string, e.g., "user.conversation.history").
• type_name: The original type name of the data (a string, e.g., "list"), used for subsequent deserialization.

Interface Design

MemoryRef

MemoryRef itself is serializable. It defines the data structure for the reference and its associated operations.

class MemoryRef(BaseModel):
    """
    A serializable, persistent reference to a specific data item in Short-Term Memory.
    """
    path: str
    type_name: str

    @staticmethod
    def create(path: str, type_name: str) -> MemoryRef:
        """Create a new MemoryRef instance based on the given path and type name.
  
        Parameters
        ----------
        path : str
            The absolute path to the data in Short-Term Memory.
        type_name : str
            The original Python type name of the data.
  
        Returns:
        ----------
        MemoryRef
            A new MemoryRef instance.
        """
    
    def resolve(self, ctx: RunnerContext) -> Any:
        """Resolve the reference to retrieve the actual data.
        
        Parameters
        ----------
        ctx : RunnerContext
            The current execution context, used to access Short-Term Memory.
        
        Returns:
        -------
        Any
            The deserialized, original data.
        """

Extension to MemoryObject

To make operations more flexible and convenient, we extend the core methods of MemoryObject to support the creation and resolution of MemoryRef.

class MemoryObject(ABC):

    @abstractmethod
    def set(self, path: str, value: Any) -> MemoryRef:
        """Set data to a specified path and returns a corresponding reference.

        Parameters
        ----------
        path : str
            The dot-separated path used to locate the data.
        value : Any
            Any serializable object to be stored.

        Returns:
        -------
        MemoryRef
            A newly created MemoryRef instance pointing to the just-stored data.
        """

    @abstractmethod
    def get(self, path_or_ref: Union[str, MemoryRef]) -> Any:
        """Retrieve data from a specified path or by a given reference.

        Parameters
        ----------
        path_or_ref : Union[str, MemoryRef]
            The path or reference to the data. 
              If a string path is provided, return the value stored at that location. 
              If a MemoryRef is provided, return the deserialized original data.

        Returns:
        -------
        Any
            The retrieved data.
        """

Agent Example

Scenario Description: This example involves two Actions. Action A receives an InputEvent, stores a large data object in Short-Term Memory, and passes the corresponding reference. Action B then receives this reference, resolves it to get the actual data, processes it, and outputs the final result.

class CustomDataSet(BaseModel):
    """Custom data object."""
    id: str
    data: bytes

class MyEvent(Event):
    """
    A custom Event with a payload field that can hold either the actual data or a reference.
    """
    payload: Union[CustomDataSet, MemoryRef]

@action(InputEvent)
def action_a(event: InputEvent, ctx: RunnerContext):
    """
    Store the data in Short-Term Memory and pass its reference.
    """
    # Assume the input is a dictionary, such as: {'id': 'large_002', 'content': '...'}
    input_dict = event.input
    large_data_content = input_dict['content'].encode('utf-8')
    data_to_store = CustomDataSet(id=input_dict['id'], data=large_data_content)
    
    stm = ctx.get_short_term_memory()
    data_path = "shared_datasets." + input_dict['id'] + ".content"
  
    # Approach 1: Store data, then explicitly create the reference.
    stm.set(data_path, data_to_store)
    data_ref = MemoryRef.create(data_path, "CustomDataSet")
    # Approach 2: Store data and retrieve the returned reference directly.
    data_ref = stm.set(data_path, data_to_store)
    
    ctx.send_event(MyEvent(payload=data_ref))
  
@action(MyEvent)
def action_b(event: MyEvent, ctx: RunnerContext):
    """
    Receive a field and process it to get the actual data.
    """
    payload = event.payload
    
    # Approach 1: Manually check the type of the payload and use the MemoryRef.resolve() method.
    if isinstance(payload, MemoryRef):
        resolved_data = payload.resolve(ctx)
    else:
        resolved_data = payload
    # Approach 2: Use stm.get() to resolve the MemoryRef automatically.
    stm = ctx.get_short_term_memory()
    resolved_data = stm.get(payload)
  
    data_size = len(resolved_data.data)
  
    output_message = f"Processed data with ID '{resolved_data.id}', size: {data_size} bytes."
    ctx.send_event(OutputEvent(output=output_message))