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

### 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))

```

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

----
This is an automatically sent email for [email protected].
To unsubscribe, please send an email to: [email protected]

Reply via email to