agentscope-runtime 0.1.5b2__py3-none-any.whl → 0.2.0b1__py3-none-any.whl

agentscope_runtime/engine/services/sandbox_service.py

@@ -2,27 +2,17 @@
 # pylint: disable=too-many-branches
 from typing import List
 
-try:
-    from ...sandbox.enums import SandboxType
-    from ...sandbox.manager import SandboxManager
-    from ...sandbox.registry import SandboxRegistry
-    from ...sandbox.tools.mcp_tool import MCPTool
-    from ...sandbox.tools.sandbox_tool import SandboxTool
-    from ...sandbox.tools.function_tool import FunctionTool
-except ImportError:
-    pass
-
+from ...sandbox.enums import SandboxType
+from ...sandbox.manager import SandboxManager
+from ...sandbox.registry import SandboxRegistry
+from ...sandbox.tools.mcp_tool import MCPTool
+from ...sandbox.tools.sandbox_tool import SandboxTool
+from ...sandbox.tools.function_tool import FunctionTool
 from ...engine.services.base import ServiceWithLifecycleManager
 
 
 class SandboxService(ServiceWithLifecycleManager):
     def __init__(self, base_url=None, bearer_token=None):
-        if SandboxManager is None:
-            raise ImportError(
-                "SandboxManager is not available. "
-                "Please install agentscope-runtime[sandbox]",
-            )
-
         self.manager_api = SandboxManager(
             base_url=base_url,
             bearer_token=bearer_token,
@@ -30,18 +20,20 @@ class SandboxService(ServiceWithLifecycleManager):
 
         self.base_url = base_url
         self.bearer_token = bearer_token
-        self.sandbox_type_set = set(item.value for item in SandboxType)
-
-        self.session_mapping = {}
 
     async def start(self) -> None:
         pass
 
     async def stop(self) -> None:
         # Release all environments
-        for _, env_ids in self.session_mapping.items():
-            for env_id in env_ids:
-                self.manager_api.release(env_id)
+        session_keys = self.manager_api.list_session_keys()
+
+        if session_keys:
+            for session_ctx_id in session_keys:
+                env_ids = self.manager_api.get_session_mapping(session_ctx_id)
+                if env_ids:
+                    for env_id in env_ids:
+                        self.manager_api.release(env_id)
 
         if self.base_url is None:
             # Embedded mode
@@ -57,61 +49,52 @@ class SandboxService(ServiceWithLifecycleManager):
         env_types=None,
         tools=None,
     ) -> List:
-        for tool in tools:
-            if not isinstance(tool, (SandboxTool, FunctionTool, MCPTool)):
-                raise ValueError(
-                    "tools must be instances of SandboxTool, FunctionTool, "
-                    "or MCPTool",
-                )
-
         # Create a composite key
-        composite_key = self._create_composite_key(session_id, user_id)
+        session_ctx_id = self._create_session_ctx_id(session_id, user_id)
 
-        # Check if the composite_key already has an environment
-        if composite_key in self.session_mapping:
+        env_ids = self.manager_api.get_session_mapping(session_ctx_id)
+
+        # Check if the session_ctx_id already has an environment
+        if env_ids:
             # Connect to existing environment
-            return self._connect_existing_environment(composite_key)
+            return self._connect_existing_environment(env_ids)
         else:
             # Create a new environment
             return self._create_new_environment(
-                composite_key,
+                session_ctx_id,
                 env_types,
                 tools,
             )
 
     def _create_new_environment(
         self,
-        composite_key,
+        session_ctx_id: str,
         env_types=None,
         tools=None,
     ):
+        if tools:
+            for tool in tools:
+                if not isinstance(tool, (SandboxTool, FunctionTool, MCPTool)):
+                    raise ValueError(
+                        "tools must be instances of SandboxTool, "
+                        "FunctionTool, or MCPTool",
+                    )
+
         if env_types is None:
             assert (
                 tools is not None
             ), "tools must be specified when env_types is not set"
 
-        server_configs = None
         if tools:
-            server_config_list = []
             tool_env_types = set()
             for tool in tools:
                 tool_env_types.add(tool.sandbox_type)
-                if isinstance(tool, MCPTool):
-                    server_config_list.append(tool.server_configs)
-
             if env_types is None:
                 env_types = []
 
-            if server_config_list:
-                server_configs = {"mcpServers": {}}
-                for server_config in server_config_list:
-                    server_configs["mcpServers"].update(
-                        server_config["mcpServers"],
-                    )
-
             env_types = set(env_types) | tool_env_types
 
-        wbs = []
+        sandboxes = []
         for env_type in env_types:
             if env_type is None:
                 continue
@@ -120,6 +103,7 @@ class SandboxService(ServiceWithLifecycleManager):
 
             box_id = self.manager_api.create_from_pool(
                 sandbox_type=box_type.value,
+                meta={"session_ctx_id": session_ctx_id},
             )
             box_cls = SandboxRegistry.get_classes_by_type(box_type)
 
@@ -135,22 +119,30 @@ class SandboxService(ServiceWithLifecycleManager):
                 # Embedded mode
                 box.manager_api = self.manager_api
 
-            if box_type == SandboxType.BASE:
-                if server_configs:
-                    box.add_mcp_servers(server_configs, overwrite=False)
-
-            # Store mapping from composite_key to env_id
-            if composite_key not in self.session_mapping:
-                self.session_mapping[composite_key] = []
-
-            self.session_mapping[composite_key].append(box_id)
-
-            wbs.append(box)
-        return wbs
-
-    def _connect_existing_environment(self, composite_key):
-        env_ids = self.session_mapping.get(composite_key)
-
+            # Add MCP to the sandbox
+            server_config_list = []
+            if tools:
+                for tool in tools:
+                    if isinstance(tool, MCPTool) and SandboxType(
+                        tool.sandbox_type,
+                    ) == SandboxType(box_type):
+                        server_config_list.append(tool.server_configs)
+            if server_config_list:
+                server_configs = {"mcpServers": {}}
+                for server_config in server_config_list:
+                    if (
+                        server_config is not None
+                        and "mcpServers" in server_config
+                    ):
+                        server_configs["mcpServers"].update(
+                            server_config["mcpServers"],
+                        )
+                box.add_mcp_servers(server_configs, overwrite=False)
+
+            sandboxes.append(box)
+        return sandboxes
+
+    def _connect_existing_environment(self, env_ids: List[str]):
         boxes = []
         for env_id in env_ids:
             info = self.manager_api.get_info(env_id)
@@ -183,16 +175,16 @@ class SandboxService(ServiceWithLifecycleManager):
         return boxes
 
     def release(self, session_id, user_id):
-        composite_key = self._create_composite_key(session_id, user_id)
+        session_ctx_id = self._create_session_ctx_id(session_id, user_id)
 
-        # Retrieve and remove env_id using composite_key
-        env_ids = self.session_mapping.pop(composite_key, [])
+        env_ids = self.manager_api.get_session_mapping(session_ctx_id)
 
-        for env_id in env_ids:
-            self.manager_api.release(env_id)
+        if env_ids:
+            for env_id in env_ids:
+                self.manager_api.release(env_id)
 
         return True
 
-    def _create_composite_key(self, session_id, user_id):
+    def _create_session_ctx_id(self, session_id, user_id):
         # Create a composite key from session_id and user_id
         return f"{session_id}_{user_id}"

agentscope_runtime/engine/services/tablestore_memory_service.py

@@ -0,0 +1,304 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=redefined-outer-name
+import asyncio
+import copy
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+try:
+    import tablestore
+    from langchain_community.embeddings import DashScopeEmbeddings
+    from langchain_core.embeddings import Embeddings
+    from tablestore import AsyncOTSClient as AsyncTablestoreClient
+    from tablestore import VectorMetricType
+    from tablestore_for_agent_memory.base.filter import Filters
+    from tablestore_for_agent_memory.knowledge.async_knowledge_store import (
+        AsyncKnowledgeStore,
+    )
+except ImportError as e:
+    raise ImportError(
+        "aliyun_tablestore is not available. "
+        "Please run pip install agentscope-runtime[aliyun_tablestore_ext]",
+    ) from e
+
+from ..schemas.agent_schemas import Message, MessageType
+from .memory_service import MemoryService
+from .utils.tablestore_service_utils import (
+    convert_messages_to_tablestore_documents,
+    convert_tablestore_document_to_message,
+    get_message_metadata_names,
+    tablestore_log,
+)
+
+
+class SearchStrategy(Enum):
+    FULL_TEXT = "full_text"
+    VECTOR = "vector"
+
+
+class TablestoreMemoryService(MemoryService):
+    """
+    A Tablestore-based implementation of the memory service.
+    based on tablestore_for_agent_memory
+    (https://github.com/aliyun/
+    alibabacloud-tablestore-for-agent-memory/blob/main/python/docs/knowledge_store_tutorial.ipynb).
+    """
+
+    _SEARCH_INDEX_NAME = "agentscope_runtime_knowledge_search_index_name"
+    _DEFAULT_SESSION_ID = "default"
+
+    def __init__(
+        self,
+        tablestore_client: AsyncTablestoreClient,
+        search_strategy: SearchStrategy = SearchStrategy.FULL_TEXT,
+        embedding_model: Optional[Embeddings] = DashScopeEmbeddings(),
+        vector_dimension: int = 1536,
+        table_name: Optional[str] = "agentscope_runtime_memory",
+        search_index_schema: Optional[List[tablestore.FieldSchema]] = (
+            tablestore.FieldSchema("user_id", tablestore.FieldType.KEYWORD),
+            tablestore.FieldSchema("session_id", tablestore.FieldType.KEYWORD),
+        ),
+        text_field: Optional[str] = "text",
+        embedding_field: Optional[str] = "embedding",
+        vector_metric_type: VectorMetricType = VectorMetricType.VM_COSINE,
+        **kwargs: Any,
+    ):
+        self._search_strategy = search_strategy
+        self._embedding_model = (
+            embedding_model  # the parameter is None, don't store vector.
+        )
+
+        if (
+            self._search_strategy == SearchStrategy.VECTOR
+            and self._embedding_model is None
+        ):
+            raise ValueError(
+                "Embedding model is required when search strategy is VECTOR.",
+            )
+
+        self._tablestore_client = tablestore_client
+        self._vector_dimension = vector_dimension
+        self._table_name = table_name
+        self._search_index_schema = (
+            list(search_index_schema)
+            if search_index_schema is not None
+            else None
+        )
+        self._text_field = text_field
+        self._embedding_field = embedding_field
+        self._vector_metric_type = vector_metric_type
+        self._knowledge_store: Optional[AsyncKnowledgeStore] = None
+        self._knowledge_store_init_parameter_kwargs = kwargs
+
+    async def _init_knowledge_store(self) -> None:
+        self._knowledge_store = AsyncKnowledgeStore(
+            tablestore_client=self._tablestore_client,
+            vector_dimension=self._vector_dimension,
+            enable_multi_tenant=False,
+            # enable multi tenant will make user be confused,
+            # we unify the usage of session id and user id,
+            # and allow users to configure the index themselves.
+            table_name=self._table_name,
+            search_index_name=TablestoreMemoryService._SEARCH_INDEX_NAME,
+            search_index_schema=copy.deepcopy(self._search_index_schema),
+            text_field=self._text_field,
+            embedding_field=self._embedding_field,
+            vector_metric_type=self._vector_metric_type,
+            **self._knowledge_store_init_parameter_kwargs,
+        )
+
+        await self._knowledge_store.init_table()
+
+    async def start(self) -> None:
+        """Start the tablestore service"""
+        if self._knowledge_store:
+            return
+        await self._init_knowledge_store()
+
+    async def stop(self) -> None:
+        """Close the tablestore service"""
+        if self._knowledge_store is None:
+            return
+        knowledge_store = self._knowledge_store
+        self._knowledge_store = None
+        await knowledge_store.close()
+
+    async def health(self) -> bool:
+        """Checks the health of the service."""
+        if self._knowledge_store is None:
+            tablestore_log("Tablestore memory service is not started.")
+            return False
+
+        try:
+            async for _ in await self._knowledge_store.get_all_documents():
+                return True
+            return True
+        except Exception as e:
+            tablestore_log(
+                f"Tablestore memory service "
+                f"cannot access Tablestore, error: {str(e)}.",
+            )
+            return False
+
+    async def add_memory(
+        self,
+        user_id: str,
+        messages: list,
+        session_id: Optional[str] = None,
+    ) -> None:
+        if not session_id:
+            session_id = TablestoreMemoryService._DEFAULT_SESSION_ID
+
+        if not messages:
+            return
+
+        tablestore_documents = convert_messages_to_tablestore_documents(
+            messages,
+            user_id,
+            session_id,
+            self._embedding_model,
+        )
+
+        put_tasks = [
+            self._knowledge_store.put_document(tablestore_document)
+            for tablestore_document in tablestore_documents
+        ]
+        await asyncio.gather(*put_tasks)
+
+    @staticmethod
+    async def get_query_text(message: Message) -> str:
+        if not message or message.type != MessageType.MESSAGE:
+            return ""
+
+        for content in message.content:
+            if content.type == "text":
+                return content.text
+
+        return ""
+
+    async def search_memory(
+        self,
+        user_id: str,
+        messages: list,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> list:
+        if (
+            not messages
+            or not isinstance(messages, list)
+            or len(messages) == 0
+        ):
+            return []
+
+        query = await TablestoreMemoryService.get_query_text(messages[-1])
+        if not query:
+            return []
+
+        top_k = 100
+        if (
+            filters
+            and "top_k" in filters
+            and isinstance(filters["top_k"], int)
+        ):
+            top_k = filters["top_k"]
+
+        if self._search_strategy == SearchStrategy.FULL_TEXT:
+            matched_messages = [
+                convert_tablestore_document_to_message(hit.document)
+                for hit in (
+                    await self._knowledge_store.full_text_search(
+                        query=query,
+                        metadata_filter=Filters.eq("user_id", user_id),
+                        limit=top_k,
+                        meta_data_to_get=get_message_metadata_names(),
+                    )
+                ).hits
+            ]
+        elif self._search_strategy == SearchStrategy.VECTOR:
+            matched_messages = [
+                convert_tablestore_document_to_message(hit.document)
+                for hit in (
+                    await self._knowledge_store.vector_search(
+                        query_vector=self._embedding_model.embed_query(query),
+                        metadata_filter=Filters.eq("user_id", user_id),
+                        top_k=top_k,
+                        meta_data_to_get=get_message_metadata_names(),
+                    )
+                ).hits
+            ]
+        else:
+            raise ValueError(
+                f"Unsupported search strategy: {self._search_strategy}",
+            )
+
+        return matched_messages
+
+    async def list_memory(
+        self,
+        user_id: str,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> list:
+        page_num = filters.get("page_num", 1) if filters else 1
+        page_size = filters.get("page_size", 10) if filters else 10
+
+        if page_num < 1 or page_size < 1:
+            raise ValueError("page_num and page_size must be greater than 0.")
+
+        next_token = None
+        for _ in range(page_num - 1):
+            next_token = (
+                await self._knowledge_store.search_documents(
+                    metadata_filter=Filters.eq("user_id", user_id),
+                    limit=page_size,
+                    next_token=next_token,
+                )
+            ).next_token
+            if not next_token:
+                tablestore_log(
+                    "Page number exceeds the total number of pages, "
+                    "return empty list.",
+                )
+                return []
+
+        messages = [
+            convert_tablestore_document_to_message(hit.document)
+            for hit in (
+                await self._knowledge_store.search_documents(
+                    metadata_filter=Filters.eq("user_id", user_id),
+                    limit=page_size,
+                    next_token=next_token,
+                    meta_data_to_get=get_message_metadata_names(),
+                )
+            ).hits
+        ]
+
+        return messages
+
+    async def delete_memory(
+        self,
+        user_id: str,
+        session_id: Optional[str] = None,
+    ) -> None:
+        delete_tablestore_documents = [
+            hit.document
+            for hit in (
+                await self._knowledge_store.search_documents(
+                    metadata_filter=(
+                        Filters.eq("user_id", user_id)
+                        if not session_id
+                        else Filters.logical_and(
+                            [
+                                Filters.eq("user_id", user_id),
+                                Filters.eq("session_id", session_id),
+                            ],
+                        )
+                    ),
+                )
+            ).hits
+        ]
+        delete_tasks = [
+            self._knowledge_store.delete_document(
+                tablestore_document.document_id,
+            )
+            for tablestore_document in delete_tablestore_documents
+        ]
+        await asyncio.gather(*delete_tasks)

agentscope_runtime/engine/services/tablestore_rag_service.py

@@ -0,0 +1,143 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=redefined-outer-name
+import asyncio
+import uuid
+from typing import Any, List, Optional, Union
+
+try:
+    from langchain_community.embeddings import DashScopeEmbeddings
+    from langchain_core.documents import Document
+    from langchain_core.embeddings import Embeddings
+    from tablestore import AsyncOTSClient as AsyncTablestoreClient
+    from tablestore import VectorMetricType
+    from tablestore_for_agent_memory.base.base_knowledge_store import (
+        Document as TablestoreDocument,
+    )
+    from tablestore_for_agent_memory.knowledge.async_knowledge_store import (
+        AsyncKnowledgeStore,
+    )
+except ImportError as e:
+    raise ImportError(
+        "aliyun_tablestore is not available. "
+        "Please run pip install agentscope-runtime[aliyun_tablestore_ext]",
+    ) from e
+
+from .rag_service import RAGService
+from .utils.tablestore_service_utils import tablestore_log
+
+
+class TablestoreRAGService(RAGService):
+    """
+    RAG Service using Tablestore(aliyun tablestore)
+    based on tablestore_for_agent_memory
+    (https://github.com/aliyun/
+    alibabacloud-tablestore-for-agent-memory/blob/main/python/docs/knowledge_store_tutorial.ipynb).
+    """
+
+    _SEARCH_INDEX_NAME = "agentscope_runtime_knowledge_search_index_name"
+    _DEFAULT_SESSION_ID = "default"
+
+    def __init__(
+        self,
+        tablestore_client: AsyncTablestoreClient,
+        embedding_model: Optional[Embeddings] = None,
+        vector_dimension: int = 1536,
+        table_name: Optional[str] = "agentscope_runtime_rag",
+        text_field: Optional[str] = "text",
+        embedding_field: Optional[str] = "embedding",
+        vector_metric_type: VectorMetricType = VectorMetricType.VM_COSINE,
+        **kwargs: Any,
+    ):
+        self._embedding_model = (
+            embedding_model if embedding_model else DashScopeEmbeddings()
+        )
+
+        self._tablestore_client = tablestore_client
+        self._vector_dimension = vector_dimension
+        self._table_name = table_name
+        self._text_field = text_field
+        self._embedding_field = embedding_field
+        self._vector_metric_type = vector_metric_type
+        self._knowledge_store: Optional[AsyncKnowledgeStore] = None
+        self._knowledge_store_init_parameter_kwargs = kwargs
+
+    async def _init_knowledge_store(self) -> None:
+        self._knowledge_store = AsyncKnowledgeStore(
+            tablestore_client=self._tablestore_client,
+            vector_dimension=self._vector_dimension,
+            enable_multi_tenant=False,
+            table_name=self._table_name,
+            search_index_name=TablestoreRAGService._SEARCH_INDEX_NAME,
+            text_field=self._text_field,
+            embedding_field=self._embedding_field,
+            vector_metric_type=self._vector_metric_type,
+            **self._knowledge_store_init_parameter_kwargs,
+        )
+
+        await self._knowledge_store.init_table()
+
+    async def start(self) -> None:
+        """Start the tablestore service"""
+        if self._knowledge_store:
+            return
+        await self._init_knowledge_store()
+
+    async def stop(self) -> None:
+        """Close the tablestore service"""
+        if self._knowledge_store is None:
+            return
+        knowledge_store = self._knowledge_store
+        self._knowledge_store = None
+        await knowledge_store.close()
+
+    async def health(self) -> bool:
+        """Checks the health of the service."""
+        if self._knowledge_store is None:
+            tablestore_log("Tablestore rag service is not started.")
+            return False
+
+        try:
+            async for _ in await self._knowledge_store.get_all_documents():
+                return True
+            return True
+        except Exception as e:
+            tablestore_log(
+                f"Tablestore rag service "
+                f"cannot access Tablestore, error: {str(e)}.",
+            )
+            return False
+
+    async def add_docs(self, docs: Union[Document, List[Document]]):
+        if not isinstance(docs, List):
+            docs = [docs]
+
+        contents = [doc.page_content for doc in docs]
+        # Encode in batches to reduce time consumption.
+        embeddings = self._embedding_model.embed_documents(contents)
+
+        put_tasks = [
+            # The conversion logic here is simple,
+            # so no separate conversion function is defined.
+            self._knowledge_store.put_document(
+                document=TablestoreDocument(
+                    document_id=f"document_{uuid.uuid4()}",
+                    text=doc.page_content,
+                    embedding=embedding,
+                    metadata=doc.metadata,
+                ),
+            )
+            for doc, embedding in zip(docs, embeddings)
+        ]
+        await asyncio.gather(*put_tasks)
+
+    async def retrieve(self, query: str, k: int = 1) -> list[str]:
+        matched_text = [
+            hit.document.text
+            for hit in (
+                await self._knowledge_store.vector_search(
+                    query_vector=self._embedding_model.embed_query(query),
+                    top_k=k,
+                )
+            ).hits
+        ]
+        return matched_text