agentscope-runtime 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

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] = None,
+        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: None)
+            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] = None,
+        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: None)
+            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

agentscope_runtime/engine/tracing/wrapper.py

@@ -26,7 +26,11 @@ from typing import (
 from pydantic import BaseModel
 from opentelemetry.propagate import extract
 from opentelemetry.context import attach
-from opentelemetry.trace import StatusCode, NoOpTracerProvider
+from opentelemetry.trace import (
+    ProxyTracerProvider,
+    StatusCode,
+    NoOpTracerProvider,
+)
 from opentelemetry import trace as ot_trace
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
     OTLPSpanExporter as OTLPSpanGrpcExporter,
@@ -917,12 +921,22 @@ def _get_ot_tracer() -> ot_trace.Tracer:
         ot_trace.Tracer: The OpenTelemetry tracer instance.
     """
 
-    def _get_ot_tracer_inner() -> ot_trace.Tracer:
+    def _has_existing_trace_provider() -> bool:
+        from opentelemetry.trace import _TRACER_PROVIDER
+
         existing_provider = ot_trace.get_tracer_provider()
+        if isinstance(existing_provider, NoOpTracerProvider):
+            return False
+        elif isinstance(existing_provider, ProxyTracerProvider):
+            # ProxyTracerProvider will use the _TRACER_PROVIDER as real tracer
+            # provider to get the tracer
+            return bool(_TRACER_PROVIDER)
 
-        if not isinstance(existing_provider, NoOpTracerProvider):
-            return ot_trace.get_tracer("agentscope_runtime")
+        return True
 
+    def _get_ot_tracer_inner() -> ot_trace.Tracer:
+        if _has_existing_trace_provider():
+            return ot_trace.get_tracer("agentscope_runtime")
         resource = Resource(
             attributes={
                 SERVICE_NAME: _get_service_name(),

agentscope_runtime/sandbox/__init__.py

@@ -3,22 +3,30 @@
 # This ensures SandboxRegistry.register() runs at import time.
 # Without this, lazy loading delays module import and types may not be
 # registered.
-from .box.base.base_sandbox import BaseSandbox
-from .box.browser.browser_sandbox import BrowserSandbox
-from .box.filesystem.filesystem_sandbox import FilesystemSandbox
-from .box.gui.gui_sandbox import GuiSandbox
+from .box.base.base_sandbox import BaseSandbox, BaseSandboxAsync
+from .box.browser.browser_sandbox import BrowserSandbox, BrowserSandboxAsync
+from .box.filesystem.filesystem_sandbox import (
+    FilesystemSandbox,
+    FilesystemSandboxAsync,
+)
+from .box.gui.gui_sandbox import GuiSandbox, GuiSandboxAsync
+from .box.mobile.mobile_sandbox import MobileSandbox, MobileSandboxAsync
 from .box.training_box.training_box import TrainingSandbox
 from .box.cloud.cloud_sandbox import CloudSandbox
-from .box.mobile.mobile_sandbox import MobileSandbox
 from .box.agentbay.agentbay_sandbox import AgentbaySandbox
 
 __all__ = [
     "BaseSandbox",
+    "BaseSandboxAsync",
     "BrowserSandbox",
+    "BrowserSandboxAsync",
     "FilesystemSandbox",
+    "FilesystemSandboxAsync",
     "GuiSandbox",
+    "GuiSandboxAsync",
+    "MobileSandbox",
+    "MobileSandboxAsync",
     "TrainingSandbox",
     "CloudSandbox",
-    "MobileSandbox",
     "AgentbaySandbox",
 ]

agentscope_runtime/sandbox/box/base/__init__.py

@@ -1,4 +1,4 @@
 # -*- coding: utf-8 -*-
-from .base_sandbox import BaseSandbox
+from .base_sandbox import BaseSandbox, BaseSandboxAsync
 
-__all__ = ["BaseSandbox"]
+__all__ = ["BaseSandbox", "BaseSandboxAsync"]

agentscope_runtime/sandbox/box/base/base_sandbox.py

@@ -4,7 +4,7 @@ from typing import Optional
 from ...utils import build_image_uri
 from ...registry import SandboxRegistry
 from ...enums import SandboxType
-from ...box.sandbox import Sandbox
+from ...box.sandbox import Sandbox, SandboxAsync
 from ...constant import TIMEOUT
 
 
@@ -49,3 +49,53 @@ class BaseSandbox(Sandbox):
             command (str): Shell command to execute.
         """
         return self.call_tool("run_shell_command", {"command": command})
+
+
+@SandboxRegistry.register(
+    build_image_uri("runtime-sandbox-base"),
+    sandbox_type=SandboxType.BASE_ASYNC,
+    security_level="medium",
+    timeout=TIMEOUT,
+    description="Base Sandbox (Async)",
+)
+class BaseSandboxAsync(SandboxAsync):
+    def __init__(
+        self,
+        sandbox_id: Optional[str] = None,
+        timeout: int = 3000,
+        base_url: Optional[str] = None,
+        bearer_token: Optional[str] = None,
+        sandbox_type: SandboxType = SandboxType.BASE_ASYNC,
+    ):
+        super().__init__(
+            sandbox_id,
+            timeout,
+            base_url,
+            bearer_token,
+            sandbox_type,
+        )
+
+    async def run_ipython_cell(self, code: str):
+        """
+        Run an IPython cell asynchronously.
+
+        Args:
+            code (str): IPython code to execute.
+        Returns:
+            Any: Response from sandbox execution
+        """
+        return await self.call_tool_async("run_ipython_cell", {"code": code})
+
+    async def run_shell_command(self, command: str):
+        """
+        Run a shell command asynchronously.
+
+        Args:
+            command (str): Shell command to execute.
+        Returns:
+            Any: Response from sandbox execution
+        """
+        return await self.call_tool_async(
+            "run_shell_command",
+            {"command": command},
+        )

agentscope_runtime/sandbox/box/browser/__init__.py

@@ -1,4 +1,4 @@
 # -*- coding: utf-8 -*-
-from .browser_sandbox import BrowserSandbox
+from .browser_sandbox import BrowserSandbox, BrowserSandboxAsync
 
-__all__ = ["BrowserSandbox"]
+__all__ = ["BrowserSandbox", "BrowserSandboxAsync"]