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

agentscope_runtime/engine/services/tablestore_session_history_service.py

@@ -0,0 +1,293 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=redefined-outer-name
+import asyncio
+import uuid
+from typing import Any, Dict, List, Optional, Union
+
+try:
+    import tablestore
+    from tablestore import AsyncOTSClient as AsyncTablestoreClient
+    from tablestore_for_agent_memory.base.base_memory_store import (
+        Session as TablestoreSession,
+    )
+    from tablestore_for_agent_memory.base.common import MetaType, Order
+    from tablestore_for_agent_memory.memory.async_memory_store import (
+        AsyncMemoryStore,
+    )
+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
+from .session_history_service import Session, SessionHistoryService
+from .utils.tablestore_service_utils import (
+    convert_message_to_tablestore_message,
+    convert_tablestore_session_to_session,
+    tablestore_log,
+)
+
+
+class TablestoreSessionHistoryService(SessionHistoryService):
+    """An aliyun tablestore implementation of the SessionHistoryService
+    based on tablestore_for_agent_memory
+    (https://github.com/aliyun/
+    alibabacloud-tablestore-for-agent-memory/blob/main/python/docs/memory_store_tutorial.ipynb).
+    """
+
+    _SESSION_SECONDARY_INDEX_NAME = (
+        "agentscope_runtime_session_secondary_index"
+    )
+    _SESSION_SEARCH_INDEX_NAME = "agentscope_runtime_session_search_index"
+    _MESSAGE_SECONDARY_INDEX_NAME = (
+        "agentscope_runtime_message_secondary_index"
+    )
+    _MESSAGE_SEARCH_INDEX_NAME = "agentscope_runtime_message_search_index"
+
+    def __init__(
+        self,
+        tablestore_client: AsyncTablestoreClient,
+        session_table_name: Optional[str] = "agentscope_runtime_session",
+        message_table_name: Optional[str] = "agentscope_runtime_message",
+        session_secondary_index_meta: Optional[Dict[str, MetaType]] = None,
+        session_search_index_schema: Optional[
+            List[tablestore.FieldSchema]
+        ] = None,
+        message_search_index_schema: Optional[
+            List[tablestore.FieldSchema]
+        ] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes the TablestoreSessionHistoryService."""
+        self._tablestore_client = tablestore_client
+        self._session_table_name = session_table_name
+        self._message_table_name = message_table_name
+        self._session_secondary_index_meta = session_secondary_index_meta
+        self._session_search_index_schema = session_search_index_schema
+        self._message_search_index_schema = message_search_index_schema
+        self._memory_store: Optional[AsyncMemoryStore] = None
+        self._memory_store_init_parameter_kwargs = kwargs
+
+    async def _init_memory_store(self) -> None:
+        self._memory_store = AsyncMemoryStore(
+            tablestore_client=self._tablestore_client,
+            session_table_name=self._session_table_name,
+            message_table_name=self._message_table_name,
+            session_secondary_index_name=self._SESSION_SECONDARY_INDEX_NAME,
+            session_search_index_name=self._SESSION_SEARCH_INDEX_NAME,
+            message_secondary_index_name=self._MESSAGE_SECONDARY_INDEX_NAME,
+            message_search_index_name=self._MESSAGE_SEARCH_INDEX_NAME,
+            session_secondary_index_meta=self._session_secondary_index_meta,
+            session_search_index_schema=self._session_search_index_schema,
+            message_search_index_schema=self._message_search_index_schema,
+            **self._memory_store_init_parameter_kwargs,
+        )
+
+        await self._memory_store.init_table()
+        await self._memory_store.init_search_index()
+
+    async def start(self) -> None:
+        """Start the tablestore service"""
+        if self._memory_store:
+            return
+        await self._init_memory_store()
+
+    async def stop(self) -> None:
+        """Close the tablestore service"""
+        if self._memory_store is None:
+            return
+        memory_store = self._memory_store
+        self._memory_store = None
+        await memory_store.close()
+
+    async def health(self) -> bool:
+        """Checks the health of the service."""
+        if self._memory_store is None:
+            tablestore_log(
+                "Tablestore session history service is not started.",
+            )
+            return False
+
+        try:
+            async for _ in await self._memory_store.list_all_sessions():
+                return True
+            return True
+        except Exception as e:
+            tablestore_log(
+                f"Tablestore session history service "
+                f"cannot access Tablestore, error: {str(e)}.",
+            )
+            return False
+
+    async def create_session(
+        self,
+        user_id: str,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        """Creates a new session for a given user and stores it.
+
+        Args:
+            user_id: The identifier for the user creating the session.
+            session_id: The identifier for the session to delete.
+
+        Returns:
+            A newly created Session object.
+        """
+        session_id = (
+            session_id.strip()
+            if session_id and session_id.strip()
+            else str(uuid.uuid4())
+        )
+        tablestore_session = TablestoreSession(
+            session_id=session_id,
+            user_id=user_id,
+        )
+
+        await self._memory_store.put_session(tablestore_session)
+        return convert_tablestore_session_to_session(tablestore_session)
+
+    async def get_session(
+        self,
+        user_id: str,
+        session_id: str,
+    ) -> Session | None:
+        """Retrieves a specific session from memory.
+
+        Args:
+            user_id: The identifier for the user.
+            session_id: The identifier for the session to retrieve.
+
+        Returns:
+            A Session object if found, otherwise None.
+        """
+
+        tablestore_session = await self._memory_store.get_session(
+            user_id=user_id,
+            session_id=session_id,
+        )
+
+        if not tablestore_session:
+            tablestore_session = TablestoreSession(
+                session_id=session_id,
+                user_id=user_id,
+            )
+            await self._memory_store.put_session(tablestore_session)
+            tablestore_messages = None
+        else:
+            messages_iterator = await self._memory_store.list_messages(
+                session_id=session_id,
+                order=Order.ASC,
+                # Sort by timestamp,
+                # keeping the most recent information at the end of the list.
+            )
+            tablestore_messages = [
+                message async for message in messages_iterator
+            ]
+
+        return convert_tablestore_session_to_session(
+            tablestore_session,
+            tablestore_messages,
+        )
+
+    async def delete_session(self, user_id: str, session_id: str) -> None:
+        """Deletes a specific session from memory.
+
+        If the session does not exist, the method does nothing.
+
+        Args:
+            user_id: The identifier for the user.
+            session_id: The identifier for the session to delete.
+        """
+        await self._memory_store.delete_session_and_messages(
+            user_id=user_id,
+            session_id=session_id,
+        )
+
+    async def list_sessions(self, user_id: str) -> list[Session]:
+        """Lists all sessions for a given user.
+
+        To improve performance and reduce data transfer, the returned session
+        objects do not contain the detailed response history.
+
+        Args:
+            user_id: The identifier of the user whose sessions to list.
+
+        Returns:
+            A list of Session objects belonging to the user, without history.
+        """
+        tablestore_sessions = await self._memory_store.list_sessions(user_id)
+        return [
+            convert_tablestore_session_to_session(tablestore_session)
+            async for tablestore_session in tablestore_sessions
+        ]
+
+    async def append_message(
+        self,
+        session: Session,
+        message: Union[
+            Message,
+            List[Message],
+            Dict[str, Any],
+            List[Dict[str, Any]],
+        ],
+    ) -> None:
+        """Appends message to a session's history in memory.
+
+        This method finds the authoritative session object in the in-memory
+        storage and appends the message to its history. It supports both
+        dictionary format messages and Message objects.
+
+        Args:
+            session: The session object, typically from the context. The
+                user_id and id from this object are used for lookup.
+            message: The message or list of messages to append to the
+                session's history.
+        """
+        # Normalize to list
+        if not isinstance(message, list):
+            message = [message]
+
+        norm_message = []
+        for msg in message:
+            if msg is None:
+                continue
+
+            if not isinstance(msg, Message):
+                msg = Message.model_validate(msg)
+            norm_message.append(msg)
+        session.messages.extend(norm_message)
+
+        tablestore_session = await self._memory_store.get_session(
+            session_id=session.id,
+            user_id=session.user_id,
+        )
+        if tablestore_session:
+            put_tasks = [
+                self._memory_store.put_message(
+                    convert_message_to_tablestore_message(message, session),
+                )
+                for message in norm_message
+            ]
+            await asyncio.gather(*put_tasks)
+
+        else:
+            tablestore_log(
+                f"Warning: Session {session.id} not found "
+                f"in tablestore storage for "
+                f"append_message.",
+            )
+
+    async def delete_user_sessions(self, user_id: str) -> None:
+        """
+        Deletes all session history data for a specific user.
+
+        Args:
+            user_id (str): The ID of the user whose session history data should
+             be deleted
+        """
+        delete_tasks = [
+            self.delete_session(user_id, session.id)
+            for session in (await self.list_sessions(user_id=user_id))
+        ]
+        await asyncio.gather(*delete_tasks)

agentscope_runtime/engine/services/utils/tablestore_service_utils.py

@@ -0,0 +1,352 @@
+# -*- coding: utf-8 -*-
+import copy
+import json
+from typing import Any, Dict, List, Optional, Tuple
+
+import tablestore
+from langchain_core.embeddings import Embeddings
+from tablestore import AsyncOTSClient as AsyncTablestoreClient
+from tablestore.credentials import CredentialsProvider
+from tablestore_for_agent_memory.base.base_knowledge_store import (
+    Document as TablestoreDocument,
+)
+from tablestore_for_agent_memory.base.base_memory_store import (
+    Message as TablestoreMessage,
+)
+from tablestore_for_agent_memory.base.base_memory_store import (
+    Session as TablestoreSession,
+)
+
+from ...schemas.agent_schemas import ContentType, MessageType
+from ..session_history_service import Message, Session
+
+
+def create_tablestore_client(
+    end_point: str,
+    access_key_id: str,
+    access_key_secret: str,
+    instance_name: str,
+    sts_token: Optional[str] = None,
+    region: Optional[str] = None,
+    credentials_provider: Optional[CredentialsProvider] = None,
+    retry_policy: tablestore.RetryPolicy = tablestore.WriteRetryPolicy(),
+    **kwargs: Any,
+) -> AsyncTablestoreClient:
+    return AsyncTablestoreClient(
+        end_point=end_point,
+        access_key_id=access_key_id,
+        access_key_secret=access_key_secret,
+        instance_name=instance_name,
+        region=region,
+        credentials_provider=credentials_provider,
+        sts_token=None if sts_token == "" else sts_token,
+        retry_policy=retry_policy,
+        **kwargs,
+    )
+
+
+content_list_name = "content_list"
+
+
+def exclude_None_fields_in_place(obj: Dict):
+    """Remove fields with None values from dictionary in-place"""
+    obj_copy = copy.deepcopy(obj)
+    for key, value in obj_copy.items():
+        if value is None:
+            del obj[key]
+
+
+def tablestore_log(msg: str):
+    print(msg)
+
+
+def convert_tablestore_session_to_session(
+    tablestore_session: TablestoreSession,
+    tablestore_messages: Optional[List[TablestoreMessage]] = None,
+) -> Session:
+    """Convert TablestoreSession to Session"""
+    init_json = _generate_init_json_from_tablestore_session(
+        tablestore_session,
+        tablestore_messages,
+    )
+    return Session.model_validate(init_json)
+
+
+# now, the func is not be used,
+# because the interface of session history service don't need this func,
+# just for future
+def convert_session_to_tablestore_session(
+    session: Session,
+) -> Tuple[TablestoreSession, List[TablestoreMessage]]:
+    """Convert Session to TablestoreSession and list of TablestoreMessage"""
+    tablestore_session_metadata = session.model_dump(
+        exclude={"id", "user_id", "messages"},
+    )
+    exclude_None_fields_in_place(tablestore_session_metadata)
+    tablestore_session = TablestoreSession(
+        user_id=session.user_id,
+        session_id=session.id,
+        metadata=tablestore_session_metadata,
+    )
+    tablestore_messages = [
+        convert_message_to_tablestore_message(message, session)
+        for message in session.messages
+    ]
+
+    return tablestore_session, tablestore_messages
+
+
+def convert_tablestore_message_to_message(
+    tablestore_message: TablestoreMessage,
+) -> Message:
+    """Convert TablestoreMessage to Message"""
+    init_json = _generate_init_json_from_tablestore_message(tablestore_message)
+    return Message.model_validate(init_json)
+
+
+def convert_message_to_tablestore_message(
+    message: Message,
+    session: Session,
+) -> TablestoreMessage:
+    """Convert Message to TablestoreMessage"""
+    content, content_list = _generate_tablestore_content_from_message(message)
+    tablestore_message_metadata = message.model_dump(exclude={"content", "id"})
+    tablestore_message_metadata[content_list_name] = json.dumps(
+        content_list,
+        ensure_ascii=False,
+    )
+    exclude_None_fields_in_place(tablestore_message_metadata)
+    tablestore_message = TablestoreMessage(
+        session_id=session.id,
+        message_id=message.id,
+        content=content,
+        metadata=tablestore_message_metadata,
+    )
+    return tablestore_message
+
+
+# This function is designed to
+# facilitate batch embedding computation for better performance.
+def convert_messages_to_tablestore_documents(
+    messages: List[Message],
+    user_id: str,
+    session_id: str,
+    embedding_model: Optional[Embeddings] = None,
+) -> List[TablestoreDocument]:
+    """Convert list of messages
+    to TablestoreDocuments with optional batch embedding"""
+    if not embedding_model:
+        return [
+            convert_message_to_tablestore_document(
+                message,
+                user_id,
+                session_id,
+            )
+            for message in messages
+        ]
+
+    # Batch embed messages: extract content, filter non-empty,
+    # compute embeddings, and align results with original messages
+    contents = [
+        _generate_tablestore_content_from_message(message)[0]
+        for message in messages
+    ]
+    contents_not_none = [
+        content for content in contents if content is not None
+    ]
+
+    embeddings_not_none = embedding_model.embed_documents(contents_not_none)
+    embeddings = []
+    index = 0
+    for content in contents:
+        if content is not None:
+            embeddings.append(embeddings_not_none[index])
+            index += 1
+            continue
+        embeddings.append(None)
+
+    return [
+        convert_message_to_tablestore_document(
+            message,
+            user_id,
+            session_id,
+            embedding,
+        )
+        for message, embedding in zip(messages, embeddings)
+    ]
+
+
+def convert_message_to_tablestore_document(
+    message: Message,
+    user_id: str,
+    session_id: str,
+    embedding: Optional[List[float]] = None,
+) -> TablestoreDocument:
+    """Convert Message to TablestoreDocument"""
+    content, content_list = _generate_tablestore_content_from_message(message)
+    tablestore_document_metadata = message.model_dump(
+        exclude={"content", "id"},
+    )
+    tablestore_document_metadata.update(
+        {
+            "user_id": user_id,
+            "session_id": session_id,
+            content_list_name: json.dumps(content_list, ensure_ascii=False),
+        },
+    )
+    exclude_None_fields_in_place(tablestore_document_metadata)
+    tablestore_document = TablestoreDocument(
+        document_id=message.id,
+        text=content,
+        embedding=embedding if embedding else None,
+        metadata=tablestore_document_metadata,
+    )
+    return tablestore_document
+
+
+def convert_tablestore_document_to_message(
+    tablestore_document: TablestoreDocument,
+) -> Message:
+    """Convert TablestoreDocument to Message"""
+    init_json = _generate_init_json_from_tablestore_document(
+        tablestore_document,
+    )
+    return Message.model_validate(init_json)
+
+
+def _generate_init_json_from_tablestore_session(
+    tablestore_session: TablestoreSession,
+    tablestore_messages: Optional[List[TablestoreMessage]] = None,
+) -> Dict[str, Any]:
+    """Generate initialization JSON from TablestoreSession"""
+    init_json = {
+        "id": tablestore_session.session_id,
+        "user_id": tablestore_session.user_id,
+        "messages": (
+            [
+                convert_tablestore_message_to_message(tablestore_message)
+                for tablestore_message in tablestore_messages
+            ]
+            if tablestore_messages is not None
+            else []
+        ),
+    }
+    # for fit future, having more fields in Session
+    init_json.update(tablestore_session.metadata)
+    return init_json
+
+
+def _generate_init_json_from_tablestore_message(
+    tablestore_message: TablestoreMessage,
+) -> Dict[str, Any]:
+    """Generate initialization JSON from TablestoreMessage"""
+    tablestore_message = copy.deepcopy(tablestore_message)
+    tablestore_message_content_list = tablestore_message.metadata.pop(
+        content_list_name,
+        None,
+    )
+    init_json = {
+        "id": tablestore_message.message_id,
+        "content": _generate_content_from_tablestore_content(
+            text=tablestore_message.content,
+            content_list=(
+                json.loads(tablestore_message_content_list)
+                if tablestore_message_content_list
+                else None
+            ),
+        ),
+    }
+    init_json.update(tablestore_message.metadata)
+    return init_json
+
+
+def _generate_init_json_from_tablestore_document(
+    tablestore_document: TablestoreDocument,
+) -> Dict[str, Any]:
+    """Generate initialization JSON from TablestoreDocument"""
+    tablestore_document = copy.deepcopy(tablestore_document)
+    tablestore_document_content_list = tablestore_document.metadata.pop(
+        content_list_name,
+        None,
+    )
+    init_json = {
+        "id": tablestore_document.document_id,
+        "content": _generate_content_from_tablestore_content(
+            text=tablestore_document.text,
+            content_list=(
+                json.loads(tablestore_document_content_list)
+                if tablestore_document_content_list
+                else None
+            ),
+        ),
+    }
+    init_json.update(tablestore_document.metadata)
+    return init_json
+
+
+def _generate_content_from_tablestore_content(
+    text: str,
+    content_list: List[Dict[str, Any]],
+) -> Optional[List[Dict[str, Any]]]:
+    """Generate final content from text and content list"""
+    content_list = copy.deepcopy(content_list)
+
+    if text is None:
+        return content_list
+
+    for content in content_list:
+        if content["type"] == ContentType.TEXT:
+            content["text"] = text
+            break
+    return content_list
+
+
+def _generate_tablestore_content_from_message(
+    message: Message,
+) -> Tuple[Optional[str], Optional[List[Dict[str, Any]]]]:
+    """Generate Tablestore content (text and content list) from Message"""
+    if message.content is None:
+        return None, None
+
+    content_json_list = [content.model_dump() for content in message.content]
+    content = None
+
+    if message.type != MessageType.MESSAGE:
+        return content, content_json_list
+
+    for content_json in content_json_list:
+        if content_json["type"] == ContentType.TEXT:
+            content = content_json.pop("text")
+            break
+
+    return content, content_json_list
+
+
+# This global variable will be cached to reduce computation time overhead
+message_metadata_names: Optional[List[str]] = None
+
+
+def get_message_metadata_names():
+    """Get list of message metadata field names"""
+    global message_metadata_names
+
+    if message_metadata_names is not None:
+        return message_metadata_names
+
+    message_metadata_names = list(Message.model_fields.keys())
+
+    message_metadata_exclude_names = ("id", "content")
+    message_metadata_extra_names = (
+        "document_id",
+        "text",
+        "user_id",
+        "session_id",
+        content_list_name,
+    )
+
+    for exclude_name in message_metadata_exclude_names:
+        message_metadata_names.remove(exclude_name)
+    for extra_name in message_metadata_extra_names:
+        message_metadata_names.append(extra_name)
+
+    return message_metadata_names

agentscope_runtime/engine/tracing/__init__.py

@@ -3,6 +3,14 @@ from typing import Any, List, Union
 
 from .base import BaseLogHandler, Tracer, TracerHandler
 from .tracing_metric import TraceType
+from .tracing_util import TracingUtil
+from .wrapper import trace
+
+__all__ = [
+    "trace",
+    "TraceType",
+    "TracingUtil",
+]
 
 
 def create_handler(
@@ -15,9 +23,7 @@ def create_handler(
     if "default" in eval_mode:
         handlers.append(BaseLogHandler())
     elif "local_logging" in eval_mode:
-        from .local_logging_handler import (
-            LocalLogHandler,
-        )
+        from .local_logging_handler import LocalLogHandler
 
         handlers.append(LocalLogHandler(**eval_params))
     return handlers

agentscope_runtime/engine/tracing/asyncio_util.py

@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+from typing import AsyncIterable, AsyncIterator, Tuple, TypeVar
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+async def aenumerate(
+    asequence: AsyncIterable[T_co],
+    start: int = 0,
+) -> AsyncIterator[Tuple[int, T_co]]:
+    """Asynchronously enumerate an async iterator from a given start value.
+
+    Args:
+        asequence (AsyncIterable[T_co]): The async iterable to enumerate.
+        start (int): The starting value for enumeration. Defaults to 0.
+
+    Yields:
+        Tuple[int, T_co]: A tuple containing the index and the item from the
+        async iterable.
+    """
+    n = start
+    async for elem in asequence:
+        yield n, elem
+        n += 1