agentscope-runtime 1.0.5.post1__py3-none-any.whl → 1.1.0b2__py3-none-any.whl

agentscope_runtime/engine/services/memory/tablestore_memory_service.py

@@ -1,301 +0,0 @@
-# -*- coding: utf-8 -*-
-# pylint: disable=redefined-outer-name
-import asyncio
-import copy
-from enum import Enum
-from typing import Any, Dict, List, Optional
-
-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,
-)
-
-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] = None,
-        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,
-    ):
-        if embedding_model is None:
-            embedding_model = DashScopeEmbeddings()
-
-        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/session_history/__init__.py

@@ -1,32 +0,0 @@
-# -*- coding: utf-8 -*-
-from typing import TYPE_CHECKING
-from ....common.utils.lazy_loader import install_lazy_loader
-from ....common.utils.deprecation import deprecated_module
-
-deprecated_module(
-    module_name=__name__,
-    removed_in="v1.1",
-    alternative="agentscope.memory",
-)
-
-if TYPE_CHECKING:
-    from .session_history_service import (
-        SessionHistoryService,
-        InMemorySessionHistoryService,
-    )
-    from .redis_session_history_service import RedisSessionHistoryService
-    from .tablestore_session_history_service import (
-        TablestoreSessionHistoryService,
-    )
-    from .session_history_service_factory import SessionHistoryServiceFactory
-
-install_lazy_loader(
-    globals(),
-    {
-        "SessionHistoryService": ".session_history_service",
-        "InMemorySessionHistoryService": ".session_history_service",
-        "RedisSessionHistoryService": ".redis_session_history_service",
-        "TablestoreSessionHistoryService": ".tablestore_session_history_service",  # noqa
-        "SessionHistoryServiceFactory": ".session_history_service_factory",
-    },
-)

agentscope_runtime/engine/services/session_history/redis_session_history_service.py

@@ -1,283 +0,0 @@
-# -*- coding: utf-8 -*-
-import uuid
-
-from typing import Optional, Dict, Any, List, Union
-
-import redis.asyncio as aioredis
-
-from .session_history_service import SessionHistoryService
-from ...schemas.session import Session
-from ...schemas.agent_schemas import Message
-
-
-class RedisSessionHistoryService(SessionHistoryService):
-    def __init__(
-        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.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"
-        except Exception:
-            return False
-
-    def _session_key(self, user_id: str, session_id: str):
-        return f"session:{user_id}:{session_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()
-
-    def _session_from_json(self, s: str) -> Session:
-        return Session.model_validate_json(s)
-
-    async def create_session(
-        self,
-        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:
-            sid = str(uuid.uuid4())
-
-        session = Session(id=sid, user_id=user_id, messages=[])
-        key = self._session_key(user_id, sid)
-
-        await self._redis.set(key, self._session_to_json(session))
-
-        # 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(
-        self,
-        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:
-            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)
-
-    async def list_sessions(self, user_id: str) -> list[Session]:
-        """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 = []
-        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(
-        self,
-        session: Session,
-        message: Union[
-            "Message",
-            List["Message"],
-            Dict[str, Any],
-            List[Dict[str, Any]],
-        ],
-    ):
-        if not self._redis:
-            raise RuntimeError("Redis connection is not available")
-        if not isinstance(message, list):
-            message = [message]
-        norm_message = []
-        for msg in message:
-            if msg is not None:
-                if not isinstance(msg, Message):
-                    msg = Message.model_validate(msg)
-                norm_message.append(msg)
-
-        session.messages.extend(norm_message)
-
-        user_id = session.user_id
-        session_id = session.id
-        key = self._session_key(user_id, session_id)
-
-        session_json = await self._redis.get(key)
-        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
-        """
-        if not self._redis:
-            raise RuntimeError("Redis connection is not available")
-
-        pattern = self._session_pattern(user_id)
-        cursor = 0
-
-        while True:
-            cursor, keys = await self._redis.scan(
-                cursor,
-                match=pattern,
-                count=100,
-            )
-            if keys:
-                await self._redis.delete(*keys)
-
-            if cursor == 0:
-                break