The persistent state management(save_state) enhanced by TTL policy

[withsmilo]

Hi, AutoGen team.

I already know about team.save_state() and team.load_state() from https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/tutorial/state.html#saving-and-loading-teams. However, if we use persistent state through external DB, all conversation history will be stored in message_thread, which will affect the storage capacity of external DB.

So I suggest adding a new TTL policy to the state management of AutoGen. Each item stored in message_thread has a timestamp, and AutoGen does not store items smaller than a specific timestamp when users call team.save_state().

Is there a better way? Thanks.

[ekzhu]

I feel this can follow once we have a persistent layer for teams, e.g., rather than using the in-memory state, we always persist state in a storage while running. TTL could be then something implemented with that persistent layer.

I feel this is beyond the prototyping stage. Curious what is your scenario for needing this type of feature?

[withsmilo]

Hi, @ekzhu .
Please imagine we're serving AutoGen on FastAPI as follows:

1. receive a user message
2. initialize team
3. load team state from external database
4. team.load_state()
5. team.run_stream(user message)
6. team.save_state()
7. save team state to external database
8. return result

For example, let's say the user already sent us two messages. The message_thread would be like this.

[0] source: user, TextMessage: hi
[1] source: TestAgent, TextMessage: blahblah
[2] source: user, TextMessage: I'm hungry
[3] source: TestAgent, type: ToolCallRequestEvent
[4] source: TestAgent, type: ToolCallExecutionEvent
[5] source: TestAgent, type: ToolCallSummaryMessage
[6] source: TestAgent, TextMessage: blahblah

However, if users keep sending multiple messages, the length of message_thread will keep increasing as well, so our external database could reach capacity limits one day. To prevent that, I thought the following TTL policy.

[0] source: user, TextMessage: hi, timestamp: 01:00:00 (hr:min:sec)
[1] source: TestAgent, TextMessage: blahblah, timestamp: 01:00:01
[2] source: user, TextMessage: I'm hungry, timestamp: 01:10:00
[3] source: TestAgent, type: ToolCallRequestEvent, timestamp: 01:10:01
[4] source: TestAgent, type: ToolCallExecutionEvent, timestamp: 01:10:02
[5] source: TestAgent, type: ToolCallSummaryMessage, timestamp: 01:10:03
[6] source: TestAgent, TextMessage: blahblah, timestamp: 01:10:04

If ttl threshold is set to 5 minutes, [0] and [1] will be deleted, and message_thread will be as follows. I think this feature can keep our external database safe from reaching capacity limits.

[2] source: user, TextMessage: I'm hungry, timestamp: 01:10:00
[3] source: TestAgent, type: ToolCallRequestEvent, timestamp: 01:10:01
[4] source: TestAgent, type: ToolCallExecutionEvent, timestamp: 01:10:02
[5] source: TestAgent, type: ToolCallSummaryMessage, timestamp: 01:10:03
[6] source: TestAgent, TextMessage: blahblah, timestamp: 01:10:04

If I misunderstood the internal implementation of AutoGen, please let me know. Thank you always.

[ekzhu]

THanks for sharing your scenario! You are right this increasing load will make the database overloaded.

I think we can introduce a simple storage abstraction for message_thread, and TTL policy could be a parameter for an implementation of this storage abstraction. What do you think?

[withsmilo]

It looks good to me. Thanks.

[ekzhu]

Created issue here: #6227

Contributions appreciated.

[withsmilo]

Hi, @ekzhu .
Direct modification of message_thread seems more complicated than expected. How about the method below to solve the problem more easily? If we add created_at to both BaseChatMessage and BaseAgentEvent classes that store the time these Pydantic model instances are generated, users will be able to use created_at to build up a customized external persisting state management layer for their case.

from datetime import datetime, timezone

class BaseChatMessage(BaseMessage, ABC):
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    ...

class BaseAgentEvent(BaseMessage, ABC):
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    ...

This is so simple that we don't have to change the internal structure of AutoGen too much. If you think this suggestion has no problem, I'll create a new PR. Thanks.

[ekzhu]

I like it. It is much simpler.

[withsmilo]

@ekzhu I just created a related PR (#6557). Thanks.