agentscope-runtime 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

agentscope_runtime/engine/services/agent_state/state_service_factory.py

@@ -43,13 +43,10 @@ class StateServiceFactory(ServiceFactory[StateService]):
 
 StateServiceFactory.register_backend(
     "in_memory",
-    lambda **kwargs: InMemoryStateService(),
+    InMemoryStateService,
 )
 
 StateServiceFactory.register_backend(
     "redis",
-    lambda **kwargs: RedisStateService(
-        redis_url=kwargs.get("redis_url", "redis://localhost:6379/0"),
-        redis_client=kwargs.get("redis_client"),
-    ),
+    RedisStateService,
 )

agentscope_runtime/engine/services/memory/redis_memory_service.py

@@ -17,23 +17,70 @@ class RedisMemoryService(MemoryService):
         self,
         redis_url: str = "redis://localhost:6379/0",
         redis_client: Optional[aioredis.Redis] = None,
+        socket_timeout: Optional[float] = 5.0,
+        socket_connect_timeout: Optional[float] = 5.0,
+        max_connections: Optional[int] = 50,
+        retry_on_timeout: bool = True,
+        ttl_seconds: Optional[int] = 3600,  # 1 hour in seconds
+        max_messages_per_session: Optional[int] = None,
+        health_check_interval: Optional[float] = 30.0,
+        socket_keepalive: bool = True,
     ):
+        """
+        Initialize RedisMemoryService.
+
+        Args:
+            redis_url: Redis connection URL
+            redis_client: Optional pre-configured Redis client
+            socket_timeout: Socket timeout in seconds (default: 5.0)
+            socket_connect_timeout: Socket connect timeout in seconds
+             (default: 5.0)
+            max_connections: Maximum number of connections in the pool
+             (default: 50)
+            retry_on_timeout: Whether to retry on timeout (default: True)
+            ttl_seconds: Time-to-live in seconds for memory data.
+            If None, data never expires (default: 3600, i.e., 1 hour)
+            max_messages_per_session: Maximum number of messages stored per
+             session_id field within a user's Redis memory hash.
+             If None, no limit (default: None)
+            health_check_interval: Interval in seconds for health checks
+             on idle connections (default: 30.0).
+                Connections idle longer than this will be checked before reuse.
+                Set to 0 to disable.
+            socket_keepalive: Enable TCP keepalive to prevent
+            silent disconnections (default: True)
+        """
         self._redis_url = redis_url
         self._redis = redis_client
         self._DEFAULT_SESSION_ID = "default"
+        self._socket_timeout = socket_timeout
+        self._socket_connect_timeout = socket_connect_timeout
+        self._max_connections = max_connections
+        self._retry_on_timeout = retry_on_timeout
+        self._ttl_seconds = ttl_seconds
+        self._max_messages_per_session = max_messages_per_session
+        self._health_check_interval = health_check_interval
+        self._socket_keepalive = socket_keepalive
 
     async def start(self) -> None:
-        """Starts the Redis connection."""
+        """Starts the Redis connection with proper timeout
+        and connection pool settings."""
         if self._redis is None:
             self._redis = aioredis.from_url(
                 self._redis_url,
                 decode_responses=True,
+                socket_timeout=self._socket_timeout,
+                socket_connect_timeout=self._socket_connect_timeout,
+                max_connections=self._max_connections,
+                retry_on_timeout=self._retry_on_timeout,
+                health_check_interval=self._health_check_interval,
+                socket_keepalive=self._socket_keepalive,
             )
 
     async def stop(self) -> None:
         """Closes the Redis connection."""
         if self._redis:
-            await self._redis.close()
+            await self._redis.aclose()
             self._redis = None
 
     async def health(self) -> bool:
@@ -73,14 +120,27 @@ class RedisMemoryService(MemoryService):
         existing_json = await self._redis.hget(key, field)
         existing_msgs = self._deserialize(existing_json)
         all_msgs = existing_msgs + messages
+
+        # Limit the number of messages per session to prevent memory issues
+        if self._max_messages_per_session is not None:
+            if len(all_msgs) > self._max_messages_per_session:
+                # Keep only the most recent messages
+                all_msgs = all_msgs[-self._max_messages_per_session :]
+
         await self._redis.hset(key, field, self._serialize(all_msgs))
 
-    async def search_memory(
+        # Set TTL for the key if configured
+        if self._ttl_seconds is not None:
+            await self._redis.expire(key, self._ttl_seconds)
+
+    async def search_memory(  # pylint: disable=too-many-branches
         self,
         user_id: str,
         messages: list,
         filters: Optional[Dict[str, Any]] = None,
     ) -> list:
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
         key = self._user_key(user_id)
         if (
             not messages
@@ -96,29 +156,52 @@ class RedisMemoryService(MemoryService):
 
         keywords = set(query.lower().split())
 
-        all_msgs = []
-        hash_keys = await self._redis.hkeys(key)
-        for session_id in hash_keys:
-            msgs_json = await self._redis.hget(key, session_id)
-            msgs = self._deserialize(msgs_json)
-            all_msgs.extend(msgs)
-
+        # Process messages in batches to avoid loading all into memory at once
         matched_messages = []
-        for msg in all_msgs:
-            candidate_content = await self.get_query_text(msg)
-            if candidate_content:
-                msg_content_lower = candidate_content.lower()
-                if any(keyword in msg_content_lower for keyword in keywords):
-                    matched_messages.append(msg)
+        hash_keys = await self._redis.hkeys(key)
 
+        # Get top_k limit early to optimize memory usage
+        top_k = None
         if (
             filters
             and "top_k" in filters
             and isinstance(filters["top_k"], int)
         ):
-            return matched_messages[-filters["top_k"] :]
+            top_k = filters["top_k"]
 
-        return matched_messages
+        # Process each session separately to reduce memory footprint
+        for session_id in hash_keys:
+            msgs_json = await self._redis.hget(key, session_id)
+            if not msgs_json:
+                continue
+            try:
+                msgs = self._deserialize(msgs_json)
+            except Exception:
+                # Skip corrupted message data
+                continue
+
+            # Match messages in this session
+            for msg in msgs:
+                candidate_content = await self.get_query_text(msg)
+                if candidate_content:
+                    msg_content_lower = candidate_content.lower()
+                    if any(
+                        keyword in msg_content_lower for keyword in keywords
+                    ):
+                        matched_messages.append(msg)
+
+        # Apply top_k filter if specified
+        if top_k is not None:
+            result = matched_messages[-top_k:]
+        else:
+            result = matched_messages
+
+        # Refresh TTL on read to extend lifetime of actively used data,
+        # if a TTL is configured and there is existing data for this key.
+        if self._ttl_seconds is not None and hash_keys:
+            await self._redis.expire(key, self._ttl_seconds)
+
+        return result
 
     async def get_query_text(self, message: Message) -> str:
         if message:
@@ -133,20 +216,39 @@ class RedisMemoryService(MemoryService):
         user_id: str,
         filters: Optional[Dict[str, Any]] = None,
     ) -> list:
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
         key = self._user_key(user_id)
-        all_msgs = []
-        hash_keys = await self._redis.hkeys(key)
-        for session_id in sorted(hash_keys):
-            msgs_json = await self._redis.hget(key, session_id)
-            msgs = self._deserialize(msgs_json)
-            all_msgs.extend(msgs)
-
         page_num = filters.get("page_num", 1) if filters else 1
         page_size = filters.get("page_size", 10) if filters else 10
 
         start_index = (page_num - 1) * page_size
         end_index = start_index + page_size
 
+        # Optimize: Calculate which sessions we need to load
+        # For simplicity, we still load all but could be optimized further
+        # to only load sessions that contain the requested page range
+        all_msgs = []
+        hash_keys = await self._redis.hkeys(key)
+        for session_id in sorted(hash_keys):
+            msgs_json = await self._redis.hget(key, session_id)
+            if msgs_json:
+                try:
+                    msgs = self._deserialize(msgs_json)
+                    all_msgs.extend(msgs)
+                except json.JSONDecodeError:
+                    # Skip corrupted message data
+                    continue
+
+                # Early exit optimization: if we've loaded enough messages
+                # to cover the requested page, we can stop (but this assumes
+                # we need all previous messages for proper ordering)
+                # For now, we keep loading all for correctness
+
+        # Refresh TTL on active use to keep memory alive,
+        # mirroring get_session behavior
+        if self._ttl_seconds is not None and hash_keys:
+            await self._redis.expire(key, self._ttl_seconds)
         return all_msgs[start_index:end_index]
 
     async def delete_memory(
@@ -154,6 +256,8 @@ class RedisMemoryService(MemoryService):
         user_id: str,
         session_id: Optional[str] = None,
     ) -> None:
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
         key = self._user_key(user_id)
         if session_id:
             await self._redis.hdel(key, session_id)

agentscope_runtime/engine/services/session_history/redis_session_history_service.py

@@ -15,23 +15,73 @@ class RedisSessionHistoryService(SessionHistoryService):
         self,
         redis_url: str = "redis://localhost:6379/0",
         redis_client: Optional[aioredis.Redis] = None,
+        socket_timeout: Optional[float] = 5.0,
+        socket_connect_timeout: Optional[float] = 5.0,
+        max_connections: Optional[int] = 50,
+        retry_on_timeout: bool = True,
+        ttl_seconds: Optional[int] = 3600,  # 1 hour in seconds
+        max_messages_per_session: Optional[int] = None,
+        health_check_interval: Optional[float] = 30.0,
+        socket_keepalive: bool = True,
     ):
+        """
+        Initialize RedisSessionHistoryService.
+
+        Args:
+            redis_url: Redis connection URL
+            redis_client: Optional pre-configured Redis client
+            socket_timeout: Socket timeout in seconds (default: 5.0)
+            socket_connect_timeout: Socket connect timeout in seconds
+            (default: 5.0)
+            max_connections: Maximum number of connections in the pool
+            (default: 50)
+            retry_on_timeout: Whether to retry on timeout (default: True)
+            ttl_seconds: Time-to-live in seconds for session data.
+            If None, data never expires (default: 3600, i.e., 1 hour)
+            max_messages_per_session: Maximum number of messages per session.
+            If None, no limit (default: None)
+            health_check_interval: Interval in seconds for health checks on
+            idle connections (default: 30.0).
+                Connections idle longer than this will be checked before reuse.
+                Set to 0 to disable.
+            socket_keepalive: Enable TCP keepalive to prevent
+            silent disconnections (default: True)
+        """
         self._redis_url = redis_url
         self._redis = redis_client
+        self._socket_timeout = socket_timeout
+        self._socket_connect_timeout = socket_connect_timeout
+        self._max_connections = max_connections
+        self._retry_on_timeout = retry_on_timeout
+        self._ttl_seconds = ttl_seconds
+        self._max_messages_per_session = max_messages_per_session
+        self._health_check_interval = health_check_interval
+        self._socket_keepalive = socket_keepalive
 
     async def start(self):
+        """Starts the Redis connection with proper timeout and connection
+        pool settings."""
         if self._redis is None:
             self._redis = aioredis.from_url(
                 self._redis_url,
                 decode_responses=True,
+                socket_timeout=self._socket_timeout,
+                socket_connect_timeout=self._socket_connect_timeout,
+                max_connections=self._max_connections,
+                retry_on_timeout=self._retry_on_timeout,
+                health_check_interval=self._health_check_interval,
+                socket_keepalive=self._socket_keepalive,
             )
 
     async def stop(self):
         if self._redis:
-            await self._redis.close()
+            await self._redis.aclose()
             self._redis = None
 
     async def health(self) -> bool:
+        """Checks the health of the service."""
+        if not self._redis:
+            return False
         try:
             pong = await self._redis.ping()
             return pong is True or pong == "PONG"
@@ -41,8 +91,9 @@ class RedisSessionHistoryService(SessionHistoryService):
     def _session_key(self, user_id: str, session_id: str):
         return f"session:{user_id}:{session_id}"
 
-    def _index_key(self, user_id: str):
-        return f"session_index:{user_id}"
+    def _session_pattern(self, user_id: str):
+        """Generate the pattern for scanning session keys for a user."""
+        return f"session:{user_id}:*"
 
     def _session_to_json(self, session: Session) -> str:
         return session.model_dump_json()
@@ -55,6 +106,8 @@ class RedisSessionHistoryService(SessionHistoryService):
         user_id: str,
         session_id: Optional[str] = None,
     ) -> Session:
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
         if session_id and session_id.strip():
             sid = session_id.strip()
         else:
@@ -64,7 +117,11 @@ class RedisSessionHistoryService(SessionHistoryService):
         key = self._session_key(user_id, sid)
 
         await self._redis.set(key, self._session_to_json(session))
-        await self._redis.sadd(self._index_key(user_id), sid)
+
+        # Set TTL for the session key if configured
+        if self._ttl_seconds is not None:
+            await self._redis.expire(key, self._ttl_seconds)
+
         return session
 
     async def get_session(
@@ -72,31 +129,63 @@ class RedisSessionHistoryService(SessionHistoryService):
         user_id: str,
         session_id: str,
     ) -> Optional[Session]:
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
         key = self._session_key(user_id, session_id)
         session_json = await self._redis.get(key)
         if session_json is None:
-            session = Session(id=session_id, user_id=user_id)
-            await self._redis.set(key, self._session_to_json(session))
-            await self._redis.sadd(self._index_key(user_id), session_id)
-            return session
-        return self._session_from_json(session_json)
+            return None
+
+        try:
+            session = self._session_from_json(session_json)
+        except Exception:
+            # Return None for corrupted session data
+            return None
+
+        # Refresh TTL when accessing the session
+        if self._ttl_seconds is not None:
+            await self._redis.expire(key, self._ttl_seconds)
+
+        return session
 
     async def delete_session(self, user_id: str, session_id: str):
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
         key = self._session_key(user_id, session_id)
         await self._redis.delete(key)
-        await self._redis.srem(self._index_key(user_id), session_id)
 
     async def list_sessions(self, user_id: str) -> list[Session]:
-        idx_key = self._index_key(user_id)
-        session_ids = await self._redis.smembers(idx_key)
+        """List all sessions for a user by scanning session keys.
+
+        Uses SCAN to find all session:{user_id}:* keys. Expired sessions
+        naturally disappear as their keys expire, avoiding stale entries.
+        """
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
+        pattern = self._session_pattern(user_id)
         sessions = []
-        for sid in session_ids:
-            key = self._session_key(user_id, sid)
-            session_json = await self._redis.get(key)
-            if session_json:
-                session = self._session_from_json(session_json)
-                session.messages = []
-                sessions.append(session)
+        cursor = 0
+
+        while True:
+            cursor, keys = await self._redis.scan(
+                cursor,
+                match=pattern,
+                count=100,
+            )
+            for key in keys:
+                session_json = await self._redis.get(key)
+                if session_json:
+                    try:
+                        session = self._session_from_json(session_json)
+                        session.messages = []
+                        sessions.append(session)
+                    except Exception:
+                        # Skip corrupted session data
+                        continue
+
+            if cursor == 0:
+                break
+
         return sessions
 
     async def append_message(
@@ -109,6 +198,8 @@ class RedisSessionHistoryService(SessionHistoryService):
             List[Dict[str, Any]],
         ],
     ):
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
         if not isinstance(message, list):
             message = [message]
         norm_message = []
@@ -125,21 +216,50 @@ class RedisSessionHistoryService(SessionHistoryService):
         key = self._session_key(user_id, session_id)
 
         session_json = await self._redis.get(key)
-        if session_json:
-            stored_session = self._session_from_json(session_json)
-            stored_session.messages.extend(norm_message)
-            await self._redis.set(key, self._session_to_json(stored_session))
-            await self._redis.sadd(self._index_key(user_id), session_id)
-        else:
-            print(
-                f"Warning: Session {session.id} not found in storage for "
-                f"append_message.",
+        if session_json is None:
+            # Session expired or not found, treat as a new session
+            # Create a new session with the current messages
+            stored_session = Session(
+                id=session_id,
+                user_id=user_id,
+                messages=norm_message.copy(),
             )
+        else:
+            try:
+                stored_session = self._session_from_json(session_json)
+                stored_session.messages.extend(norm_message)
+            except Exception:
+                # Session data corrupted, treat as a new session
+                stored_session = Session(
+                    id=session_id,
+                    user_id=user_id,
+                    messages=norm_message.copy(),
+                )
+
+        # Limit the number of messages per session to prevent memory issues
+        if self._max_messages_per_session is not None:
+            if len(stored_session.messages) > self._max_messages_per_session:
+                # Keep only the most recent messages
+                stored_session.messages = stored_session.messages[
+                    -self._max_messages_per_session :
+                ]
+                # Keep the in-memory session in sync with the stored session
+                session.messages = session.messages[
+                    -self._max_messages_per_session :
+                ]
+
+        await self._redis.set(key, self._session_to_json(stored_session))
+
+        # Set TTL for the session key if configured
+        if self._ttl_seconds is not None:
+            await self._redis.expire(key, self._ttl_seconds)
 
     async def delete_user_sessions(self, user_id: str) -> None:
         """
         Deletes all session history data for a specific user.
 
+        Uses SCAN to find all session keys for the user and deletes them.
+
         Args:
             user_id (str): The ID of the user whose session history data should
              be deleted
@@ -147,11 +267,17 @@ class RedisSessionHistoryService(SessionHistoryService):
         if not self._redis:
             raise RuntimeError("Redis connection is not available")
 
-        index_key = self._index_key(user_id)
-        session_ids = await self._redis.smembers(index_key)
+        pattern = self._session_pattern(user_id)
+        cursor = 0
 
-        for session_id in session_ids:
-            key = self._session_key(user_id, session_id)
-            await self._redis.delete(key)
+        while True:
+            cursor, keys = await self._redis.scan(
+                cursor,
+                match=pattern,
+                count=100,
+            )
+            if keys:
+                await self._redis.delete(*keys)
 
-        await self._redis.delete(index_key)
+            if cursor == 0:
+                break