Re: [D] Short-Term Memory Design [flink-agents]

["KeGu-069 (via GitHub)" <gi...@apache.org>]

GitHub user KeGu-069 added a comment to the discussion: Short-Term Memory Design

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.
MemoryRef 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.
● Core Components: 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))


GitHub link: https://github.com/apache/flink-agents/discussions/40#discussioncomment-14017572

----
This is an automatically sent email for issues@flink.apache.org.
To unsubscribe, please send an email to: issues-unsubscribe@flink.apache.org