diff --git a/metagpt/memory/memory.py b/metagpt/memory/memory.py index 5dbadf9c38..fb3c483587 100644 --- a/metagpt/memory/memory.py +++ b/metagpt/memory/memory.py @@ -5,11 +5,13 @@ @Author : alexanderwu @File : memory.py @Modified By: mashenquan, 2023-11-1. According to RFC 116: Updated the type of index key. +@Modified By: 2026-04-21. Added TTL (Time-To-Live) support for messages with automatic cleanup. """ +import asyncio from collections import defaultdict from typing import DefaultDict, Iterable, Optional, Set -from pydantic import BaseModel, Field, SerializeAsAny +from pydantic import BaseModel, Field, PrivateAttr, SerializeAsAny from metagpt.const import IGNORED_MESSAGE_ID from metagpt.schema import Message @@ -24,6 +26,24 @@ class Memory(BaseModel): index: DefaultDict[str, list[SerializeAsAny[Message]]] = Field(default_factory=lambda: defaultdict(list)) ignore_id: bool = False + _cleanup_task: Optional[asyncio.Task] = PrivateAttr(default=None) + _is_running: bool = PrivateAttr(default=False) + + def _cleanup_expired_messages(self) -> int: + """Remove all expired messages from storage and index. + + Returns: + int: The number of expired messages that were removed. + """ + expired_messages = [message for message in self.storage if message.is_expired()] + + for message in expired_messages: + self.storage.remove(message) + if message.cause_by and message in self.index[message.cause_by]: + self.index[message.cause_by].remove(message) + + return len(expired_messages) + def add(self, message: Message): """Add a new message to storage, while updating the index""" if self.ignore_id: @@ -40,10 +60,12 @@ def add_batch(self, messages: Iterable[Message]): def get_by_role(self, role: str) -> list[Message]: """Return all messages of a specified role""" + self._cleanup_expired_messages() return [message for message in self.storage if message.role == role] def get_by_content(self, content: str) -> list[Message]: """Return all messages containing a specified content""" + self._cleanup_expired_messages() return [message for message in self.storage if content in message.content] def delete_newest(self) -> "Message": @@ -71,14 +93,17 @@ def clear(self): def count(self) -> int: """Return the number of messages in storage""" + self._cleanup_expired_messages() return len(self.storage) def try_remember(self, keyword: str) -> list[Message]: """Try to recall all messages containing a specified keyword""" + self._cleanup_expired_messages() return [message for message in self.storage if keyword in message.content] def get(self, k=0) -> list[Message]: """Return the most recent k memories, return all when k=0""" + self._cleanup_expired_messages() return self.storage[-k:] def find_news(self, observed: list[Message], k=0) -> list[Message]: @@ -93,11 +118,13 @@ def find_news(self, observed: list[Message], k=0) -> list[Message]: def get_by_action(self, action) -> list[Message]: """Return all messages triggered by a specified Action""" + self._cleanup_expired_messages() index = any_to_str(action) return self.index[index] def get_by_actions(self, actions: Set) -> list[Message]: """Return all messages triggered by specified Actions""" + self._cleanup_expired_messages() rsp = [] indices = any_to_str_set(actions) for action in indices: @@ -108,5 +135,58 @@ def get_by_actions(self, actions: Set) -> list[Message]: @handle_exception def get_by_position(self, position: int) -> Optional[Message]: - """Returns the message at the given position if valid, otherwise returns None""" + """Returns the message at the given position if valid and not expired, otherwise returns None""" + self._cleanup_expired_messages() + if position < 0 or position >= len(self.storage): + return None return self.storage[position] + + async def _periodic_cleanup(self, interval: int = 60): + """Asynchronous background task that periodically cleans up expired messages. + + Args: + interval: The time in seconds between each cleanup check. + """ + while self._is_running: + try: + self._cleanup_expired_messages() + await asyncio.sleep(interval) + except asyncio.CancelledError: + break + except Exception as e: + import logging + logger = logging.getLogger(__name__) + logger.error(f"Error in periodic cleanup: {e}") + await asyncio.sleep(interval) + + def start(self, cleanup_interval: int = 60): + """Start the background periodic cleanup task. + + Args: + cleanup_interval: The time in seconds between each cleanup check. Defaults to 60. + """ + if self._is_running: + return + + self._is_running = True + + try: + loop = asyncio.get_running_loop() + self._cleanup_task = loop.create_task(self._periodic_cleanup(cleanup_interval)) + except RuntimeError: + import threading + from metagpt.utils.async_helper import run_coroutine_in_new_loop + + def run_cleanup(): + asyncio.run(self._periodic_cleanup(cleanup_interval)) + + thread = threading.Thread(target=run_cleanup, daemon=True) + thread.start() + + def stop(self): + """Stop the background periodic cleanup task gracefully.""" + self._is_running = False + + if self._cleanup_task: + self._cleanup_task.cancel() + self._cleanup_task = None diff --git a/metagpt/schema.py b/metagpt/schema.py index 52badcc21a..a82143c16b 100644 --- a/metagpt/schema.py +++ b/metagpt/schema.py @@ -240,6 +240,8 @@ class Message(BaseModel): sent_from: str = Field(default="", validate_default=True) send_to: set[str] = Field(default={MESSAGE_ROUTE_TO_ALL}, validate_default=True) metadata: Dict[str, Any] = Field(default_factory=dict) # metadata for `content` and `instruct_content` + ttl: int = Field(default=-1, validate_default=True) # Time-To-Live in seconds, -1 means never expire + created_at: float = Field(default_factory=time.time) # Creation time in seconds since epoch @field_validator("id", mode="before") @classmethod @@ -415,6 +417,18 @@ def is_user_message(self) -> bool: def is_ai_message(self) -> bool: return self.role == "assistant" + def is_expired(self) -> bool: + """Check if the message has expired based on its TTL. + + Returns: + bool: True if the message has expired, False otherwise. + Messages with ttl=-1 never expire. + """ + if self.ttl == -1: + return False + current_time = time.time() + return current_time - self.created_at > self.ttl + class UserMessage(Message): """便于支持OpenAI的消息 diff --git a/tests/metagpt/memory/test_memory.py b/tests/metagpt/memory/test_memory.py index a072b61deb..ca25b0c120 100644 --- a/tests/metagpt/memory/test_memory.py +++ b/tests/metagpt/memory/test_memory.py @@ -2,6 +2,8 @@ # -*- coding: utf-8 -*- # @Desc : the unittest of Memory +import time + from metagpt.actions import UserRequirement from metagpt.memory.memory import Memory from metagpt.schema import Message @@ -55,3 +57,268 @@ def test_memory(): memory.clear() assert memory.count() == 0 assert len(memory.index) == 0 + + +def test_message_ttl_and_created_at(): + """Test Message class ttl and created_at fields""" + # Test default values + message = Message(content="test message", role="user") + assert message.ttl == -1 + assert message.created_at > 0 + + # Test custom ttl + message_with_ttl = Message(content="test message with ttl", role="user", ttl=60) + assert message_with_ttl.ttl == 60 + + # Test is_expired method for ttl=-1 (never expire) + assert message.is_expired() == False + + # Test serialization + dumped = message.dump() + loaded = Message.load(dumped) + assert loaded.ttl == message.ttl + assert abs(loaded.created_at - message.created_at) < 0.001 + + +def test_message_expiration(): + """Test message expiration functionality""" + # Create a message that expires in 1 second + message = Message(content="expiring message", role="user", ttl=1) + assert message.is_expired() == False + + # Wait for message to expire + time.sleep(1.1) + assert message.is_expired() == True + + # Create a message that never expires + message_never_expire = Message(content="never expiring message", role="user", ttl=-1) + assert message_never_expire.is_expired() == False + + # Wait and verify it still doesn't expire + time.sleep(0.5) + assert message_never_expire.is_expired() == False + + +def test_memory_filter_expired_messages(): + """Test Memory class filtering expired messages""" + memory = Memory() + + # Create messages with different TTLs + message1 = Message(content="never expire", role="user1", ttl=-1) + message2 = Message(content="expire in 1 sec", role="user2", ttl=1) + message3 = Message(content="never expire too", role="user1", ttl=-1) + + # Add all messages to memory + memory.add_batch([message1, message2, message3]) + assert memory.count() == 3 + + # Wait for message2 to expire + time.sleep(1.1) + + # Test get() method filters expired messages + messages = memory.get() + assert len(messages) == 2 + assert message2 not in messages + + # Test get_by_role() method filters expired messages + messages = memory.get_by_role("user2") + assert len(messages) == 0 + + # Test get_by_role() for non-expired messages + messages = memory.get_by_role("user1") + assert len(messages) == 2 + + # Test get_by_content() method filters expired messages + messages = memory.get_by_content("expire") + assert len(messages) == 2 + assert message2 not in messages + + # Test try_remember() method filters expired messages + messages = memory.try_remember("expire") + assert len(messages) == 2 + assert message2 not in messages + + # Test get_by_action() method filters expired messages + messages = memory.get_by_action(UserRequirement) + assert len(messages) == 2 + assert message2 not in messages + + # Test get_by_actions() method filters expired messages + messages = memory.get_by_actions({UserRequirement}) + assert len(messages) == 2 + assert message2 not in messages + + # Test get_by_position() method returns None for expired messages + # Note: message2 is at position 1 in storage + message = memory.get_by_position(1) + assert message is None + + +def test_memory_backward_compatibility(): + """Test backward compatibility with existing code""" + memory = Memory() + + # Create messages without specifying ttl (should use default -1) + message1 = Message(content="message1", role="user1") + message2 = Message(content="message2", role="user2") + + # Verify default ttl is -1 + assert message1.ttl == -1 + assert message2.ttl == -1 + + # Add to memory + memory.add_batch([message1, message2]) + + # Verify all retrieval methods work as before + assert memory.count() == 2 + + messages = memory.get() + assert len(messages) == 2 + + messages = memory.get_by_role("user1") + assert len(messages) == 1 + assert messages[0].content == "message1" + + messages = memory.get_by_content("message") + assert len(messages) == 2 + + # Verify messages don't expire + time.sleep(0.5) + assert message1.is_expired() == False + assert message2.is_expired() == False + + messages = memory.get() + assert len(messages) == 2 + + +def test_memory_physical_cleanup(): + """Test that expired messages are physically removed from storage""" + memory = Memory() + + # Create messages with different TTLs + message1 = Message(content="never expire", role="user1", ttl=-1) + message2 = Message(content="expire in 1 sec", role="user2", ttl=1) + message3 = Message(content="never expire too", role="user1", ttl=-1) + + # Add all messages to memory + memory.add_batch([message1, message2, message3]) + assert len(memory.storage) == 3 + assert memory.count() == 3 + + # Wait for message2 to expire + time.sleep(1.1) + + # Verify message2 is still in storage before retrieval + assert len(memory.storage) == 3 + + # Call get() which should trigger cleanup + messages = memory.get() + assert len(messages) == 2 + assert message2 not in messages + + # Verify message2 is physically removed from storage + assert len(memory.storage) == 2 + assert message2 not in memory.storage + + # Verify index is also updated + index_key = message2.cause_by + if index_key in memory.index: + assert message2 not in memory.index[index_key] + + # Test count() also triggers cleanup + message4 = Message(content="expire in 0.5 sec", role="user3", ttl=0.5) + memory.add(message4) + assert len(memory.storage) == 3 + assert memory.count() == 3 + + time.sleep(0.6) + assert len(memory.storage) == 3 + + count = memory.count() + assert count == 2 + assert len(memory.storage) == 2 + + +def test_memory_cleanup_expired_messages_method(): + """Test the _cleanup_expired_messages method directly""" + memory = Memory() + + # Create messages with different TTLs + message1 = Message(content="never expire", role="user1", ttl=-1) + message2 = Message(content="expire in 0.5 sec", role="user2", ttl=0.5) + message3 = Message(content="expire in 0.5 sec too", role="user3", ttl=0.5) + + # Add all messages to memory + memory.add_batch([message1, message2, message3]) + assert len(memory.storage) == 3 + + # Wait for messages to expire + time.sleep(0.6) + + # Call cleanup method + removed_count = memory._cleanup_expired_messages() + assert removed_count == 2 + assert len(memory.storage) == 1 + assert message1 in memory.storage + assert message2 not in memory.storage + assert message3 not in memory.storage + + # Call cleanup again - should remove 0 messages + removed_count = memory._cleanup_expired_messages() + assert removed_count == 0 + + +async def test_memory_periodic_cleanup(): + """Test the periodic cleanup mechanism""" + import asyncio + + memory = Memory() + + # Create messages with different TTLs + message1 = Message(content="never expire", role="user1", ttl=-1) + message2 = Message(content="expire in 0.5 sec", role="user2", ttl=0.5) + message3 = Message(content="never expire too", role="user1", ttl=-1) + + # Add all messages to memory + memory.add_batch([message1, message2, message3]) + assert len(memory.storage) == 3 + + # Start periodic cleanup with very short interval + memory.start(cleanup_interval=0.1) + + # Wait for message to expire and cleanup to run + await asyncio.sleep(0.7) + + # Verify message2 was cleaned up + assert len(memory.storage) == 2 + assert message2 not in memory.storage + + # Stop the cleanup task + memory.stop() + + # Verify _is_running is False + assert not memory._is_running + + +def test_memory_start_stop_idempotent(): + """Test that start() and stop() are idempotent""" + memory = Memory() + + # Calling stop() before start() should not raise errors + memory.stop() + assert not memory._is_running + + # Calling start() multiple times should not create multiple tasks + memory.start(cleanup_interval=60) + assert memory._is_running + + # Calling start() again should be a no-op + memory.start(cleanup_interval=30) # Different interval but still running + assert memory._is_running + + # Calling stop() multiple times should not raise errors + memory.stop() + assert not memory._is_running + + memory.stop() + assert not memory._is_running