agentscope-runtime 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

agentscope_runtime/engine/schemas/agent_schemas.py

@@ -2,7 +2,7 @@
 # type: ignore
 from copy import deepcopy
 from datetime import datetime
-from typing import List, Dict, Optional, Any
+from typing import List, Dict, Optional, Any, Literal, TypeAlias, Annotated
 from typing import Union
 
 try:
@@ -12,6 +12,7 @@ except ImportError:
 from uuid import uuid4
 
 from pydantic import BaseModel, Field, field_validator
+from openai.types.chat import ChatCompletionChunk
 
 
 class MessageType:
@@ -29,6 +30,15 @@ class MessageType:
     HEARTBEAT = "heartbeat"
     ERROR = "error"
 
+    @classmethod
+    def all_values(cls):
+        """return all constants values in MessageType"""
+        return [
+            value
+            for name, value in vars(cls).items()
+            if not name.startswith("_") and isinstance(value, str)
+        ]
+
 
 class ContentType:
     TEXT = "text"
@@ -41,6 +51,7 @@ class Role:
     ASSISTANT = "assistant"
     USER = "user"
     SYSTEM = "system"
+    TOOL = "tool"
 
 
 class RunStatus:
@@ -70,7 +81,7 @@ class FunctionParameters(BaseModel):
 
 class FunctionTool(BaseModel):
     """
-    Model class for prompt message tool.
+    Model class for message tool.
     """
 
     name: str
@@ -89,10 +100,10 @@ class FunctionTool(BaseModel):
 
 class Tool(BaseModel):
     """
-    Model class for assistant prompt message tool call.
+    Model class for assistant message tool call.
     """
 
-    type: Optional[str] = None
+    type: Optional[str] = "function"
     """The type of the tool. Currently, only `function` is supported."""
 
     function: Optional[FunctionTool] = None
@@ -141,7 +152,7 @@ class Error(BaseModel):
 
 
 class Event(BaseModel):
-    sequence_number: Optional[str] = None
+    sequence_number: Optional[int] = None
     """sequence number of event"""
 
     object: str
@@ -213,9 +224,42 @@ class Content(Event):
     msg_id: Optional[str] = None
     """message unique id"""
 
+    @staticmethod
+    def from_chat_completion_chunk(
+        chunk: ChatCompletionChunk,
+        index: Optional[int] = None,
+    ) -> Optional[Union["TextContent", "DataContent", "ImageContent"]]:
+        if not chunk.choices:
+            return None
+
+        choice = chunk.choices[0]
+        if choice.delta.content is not None:
+            return TextContent(
+                delta=True,
+                text=choice.delta.content,
+                index=index,
+            )
+        elif choice.delta.tool_calls:
+            # TODO: support multiple tool calls output
+            tool_call = choice.delta.tool_calls[0]
+            if tool_call.function is not None:
+                return DataContent(
+                    delta=True,
+                    data={
+                        "call_id": tool_call.id,
+                        "name": tool_call.function.name,
+                        "arguments": tool_call.function.arguments,
+                    },
+                    index=index,
+                )
+            else:
+                return None
+        else:
+            return None
+
 
 class ImageContent(Content):
-    type: str = ContentType.IMAGE
+    type: Literal[ContentType.IMAGE] = ContentType.IMAGE
     """The type of the content part."""
 
     image_url: Optional[str] = None
@@ -223,7 +267,7 @@ class ImageContent(Content):
 
 
 class TextContent(Content):
-    type: str = ContentType.TEXT
+    type: Literal[ContentType.TEXT] = ContentType.TEXT
     """The type of the content part."""
 
     text: Optional[str] = None
@@ -231,13 +275,27 @@ class TextContent(Content):
 
 
 class DataContent(Content):
-    type: str = ContentType.DATA
+    type: Literal[ContentType.DATA] = ContentType.DATA
     """The type of the content part."""
 
     data: Optional[Dict] = None
     """The data content."""
 
 
+AgentRole: TypeAlias = Literal[
+    Role.ASSISTANT,
+    Role.SYSTEM,
+    Role.USER,
+    Role.TOOL,
+]
+
+
+AgentContent = Annotated[
+    Union[TextContent, ImageContent, DataContent],
+    Field(discriminator="type"),
+]
+
+
 class Message(Event):
     id: str = Field(default_factory=lambda: "msg_" + str(uuid4()))
     """message unique id"""
@@ -251,13 +309,11 @@ class Message(Event):
     status: str = RunStatus.Created
     """The status of the message. in_progress, completed, or incomplete"""
 
-    role: Optional[str] = None
+    role: Optional[AgentRole] = None
     """The role of the messages author, should be in `user`,`system`,
     'assistant'."""
 
-    content: Optional[
-        List[Union[TextContent, ImageContent, DataContent]]
-    ] = None
+    content: Optional[List[AgentContent]] = None
     """The contents of the message."""
 
     code: Optional[str] = None
@@ -266,6 +322,145 @@ class Message(Event):
     message: Optional[str] = None
     """The error message of the message."""
 
+    usage: Optional[Dict] = None
+    """response usage for output"""
+
+    @staticmethod
+    def from_openai_message(message: Union[BaseModel, dict]) -> "Message":
+        """Create a message object from an openai message."""
+
+        # in case message is a Message object
+        if isinstance(message, Message):
+            return message
+
+        # make sure operation on dict object
+        if isinstance(message, BaseModel):
+            message = message.model_dump()
+
+        # in case message is a Message format dict
+        if "type" in message and message["type"] in MessageType.all_values():
+            return Message(**message)
+
+        # handle message in openai message format
+        if message["role"] == Role.ASSISTANT and "tool_calls" in message:
+            _content_list = []
+            for tool_call in message["tool_calls"]:
+                _content = DataContent(
+                    data=FunctionCall(
+                        call_id=tool_call["id"],
+                        name=tool_call["function"]["name"],
+                        arguments=tool_call["function"]["arguments"],
+                    ).model_dump(),
+                )
+                _content_list.append(_content)
+            _message = Message(
+                type=MessageType.FUNCTION_CALL,
+                content=_content_list,
+            )
+        elif message["role"] == Role.TOOL:
+            _content = DataContent(
+                data=FunctionCallOutput(
+                    call_id=message["tool_call_id"],
+                    output=message["content"],
+                ).model_dump(),
+            )
+            _message = Message(
+                type=MessageType.FUNCTION_CALL_OUTPUT,
+                content=[_content],
+            )
+        # mainly focus on matching content
+        elif isinstance(message["content"], str):
+            _content = TextContent(text=message["content"])
+            _message = Message(
+                type=MessageType.MESSAGE,
+                role=message["role"],
+                content=[_content],
+            )
+        else:
+            _content_list = []
+            for content in message["content"]:
+                if content["type"] == "image_url":
+                    _content = ImageContent(
+                        image_url=content["image_url"]["url"],
+                    )
+                elif content["type"] == "text":
+                    _content = TextContent(text=content["text"])
+                else:
+                    _content = DataContent(data=content["text"])
+                _content_list.append(_content)
+            _message = Message(
+                type=MessageType.MESSAGE,
+                role=message["role"],
+                content=_content_list,
+            )
+        return _message
+
+    def get_text_content(self) -> Optional[str]:
+        """
+        Extract the first text content from the message.
+
+        :return:
+            First text string found in the content, or None if no text content
+        """
+        if self.content is None:
+            return None
+
+        for item in self.content:
+            if isinstance(item, TextContent):
+                return item.text
+        return None
+
+    def get_image_content(self) -> List[str]:
+        """
+        Extract all image content (URLs or base64 data) from the message.
+
+        :return:
+            List of image URLs or base64 encoded strings found in the content
+        """
+        images = []
+
+        if self.content is None:
+            return images
+
+        for item in self.content:
+            if isinstance(item, ImageContent):
+                images.append(item.image_url)
+        return images
+
+    def get_audio_content(self) -> List[str]:
+        """
+        Extract all audio content (URLs or base64 data) from the message.
+
+        :return:
+            List of audio URLs or base64 encoded strings found in the content
+        """
+        audios = []
+
+        if self.content is None:
+            return audios
+
+        for item in self.content:
+            if hasattr(item, "type"):
+                if item.type == "input_audio" and hasattr(
+                    item,
+                    "input_audio",
+                ):
+                    if hasattr(item.input_audio, "data"):
+                        audios.append(item.input_audio.data)
+                    elif hasattr(item.input_audio, "base64_data"):
+                        # Construct data URL for audio
+                        format_type = getattr(
+                            item.input_audio,
+                            "format",
+                            "mp3",
+                        )
+                        audios.append(
+                            f"data:{format_type};base64,"
+                            f"{item.input_audio.base64_data}",
+                        )
+
+        return audios
+
     def add_delta_content(
         self,
         new_content: Union[TextContent, ImageContent, DataContent],
@@ -500,7 +695,7 @@ def convert_to_openai_tool_call(function: FunctionCall):
     }
 
 
-def convert_to_openai_messages(messages: List[Message]) -> list:
+def convert_to_openai_messages(messages: List[Message]) -> List[Dict]:
     """
     Convert a generic message protocol to a model-specific protocol.
     Args:

agentscope_runtime/engine/services/context_manager.py

@@ -60,7 +60,13 @@ class ContextComposer:
                 role=Role.SYSTEM,
                 content=[TextContent(type=ContentType.TEXT, text=cooked_doc)],
             )
-            session.messages.append(message)
+            if len(session.messages) >= 1:
+                last_message = session.messages[-1]
+                session.messages.remove(last_message)
+                session.messages.append(message)
+                session.messages.append(last_message)
+            else:
+                session.messages.append(message)
 
 
 class ContextManager(ServiceManager):
@@ -145,11 +151,13 @@ async def create_context_manager(
     memory_service: MemoryService = None,
     session_history_service: SessionHistoryService = None,
     rag_service: RAGService = None,
+    context_composer_cls=ContextComposer,
 ):
     manager = ContextManager(
         memory_service=memory_service,
         session_history_service=session_history_service,
         rag_service=rag_service,
+        context_composer_cls=context_composer_cls,
     )
 
     async with manager:

agentscope_runtime/engine/services/rag_service.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-from typing import Optional
-
 from .base import ServiceWithLifecycleManager
 from ..schemas.agent_schemas import Message, MessageType
 
@@ -28,8 +26,30 @@ class RAGService(ServiceWithLifecycleManager):
         return ""
 
     async def retrieve(self, query: str, k: int = 1) -> list[str]:
+        """
+        Retrieves similar documents based on the given query.
+
+        Args:
+            query (str): The query string to search for similar documents.
+            k (int, optional): The number of similar documents to retrieve.
+            Defaults to 1.
+
+        Returns:
+            list[str]: A list of document contents that are similar to
+            the query.
+        """
         raise NotImplementedError
 
+    async def start(self) -> None:
+        """Starts the service."""
+
+    async def stop(self) -> None:
+        """Stops the service."""
+
+    async def health(self) -> bool:
+        """Checks the health of the service."""
+        return True
+
 
 DEFAULT_URI = "milvus_demo.db"
 
@@ -41,48 +61,48 @@ class LangChainRAGService(RAGService):
 
     def __init__(
         self,
-        uri: Optional[str] = None,
-        docs: Optional[list[str]] = None,
+        vectorstore=None,
+        embedding=None,
     ):
-        from langchain_community.embeddings import DashScopeEmbeddings
-        from langchain_milvus import Milvus
-
-        self.Milvus = Milvus
-        self.embeddings = DashScopeEmbeddings()
-        self.vectorstore = None
-
-        if uri:
-            self.uri = uri
-            self.from_db()
-        elif docs:
-            self.uri = DEFAULT_URI
-            self.from_docs(docs)
+        # set default embedding alg.
+        if embedding is None:
+            from langchain_community.embeddings import DashScopeEmbeddings
+
+            self.embeddings = DashScopeEmbeddings()
         else:
-            docs = []
-            self.uri = DEFAULT_URI
-            self.from_docs(docs)
-
-    def from_docs(self, docs=None):
-        if docs is None:
-            docs = []
-
-        self.vectorstore = self.Milvus.from_documents(
-            documents=docs,
-            embedding=self.embeddings,
-            connection_args={
-                "uri": self.uri,
-            },
-            drop_old=False,
-        )
-
-    def from_db(self):
-        self.vectorstore = self.Milvus(
-            embedding_function=self.embeddings,
-            connection_args={"uri": self.uri},
-            index_params={"index_type": "FLAT", "metric_type": "L2"},
-        )
+            self.embeddings = embedding
+
+        # set default vectorstore class.
+        if vectorstore is None:
+            from langchain_milvus import Milvus
+
+            self.vectorstore = Milvus.from_documents(
+                [],
+                embedding=self.embeddings,
+                connection_args={
+                    "uri": DEFAULT_URI,
+                },
+                drop_old=False,
+            )
+        else:
+            self.vectorstore = vectorstore
 
     async def retrieve(self, query: str, k: int = 1) -> list[str]:
+        """
+        Retrieves similar documents based on the given query using LangChain.
+
+        Args:
+            query (str): The query string to search for similar documents.
+            k (int, optional): The number of similar documents to retrieve.
+            Defaults to 1.
+
+        Returns:
+            list[str]: A list of document contents that are similar to the
+            query.
+
+        Raises:
+            ValueError: If the vector store is not initialized.
+        """
         if self.vectorstore is None:
             raise ValueError(
                 "Vector store not initialized. Call build_index first.",
@@ -99,3 +119,77 @@ class LangChainRAGService(RAGService):
     async def health(self) -> bool:
         """Checks the health of the service."""
         return True
+
+
+class LlamaIndexRAGService(RAGService):
+    """
+    RAG Service using LlamaIndex
+    """
+
+    def __init__(
+        self,
+        vectorstore=None,
+        embedding=None,
+    ):
+        # set default embedding alg.
+        if embedding is None:
+            from langchain_community.embeddings import DashScopeEmbeddings
+
+            self.embeddings = DashScopeEmbeddings()
+        else:
+            self.embeddings = embedding
+
+        # set default vectorstore.
+        if vectorstore is None:
+            from llama_index.core import VectorStoreIndex
+            from llama_index.core.schema import Document
+            from llama_index.vector_stores.milvus import MilvusVectorStore
+
+            # Create empty documents list for initialization
+            documents = [Document(text="")]
+
+            # Initialize Milvus vector store
+            self.vector_store = MilvusVectorStore(
+                uri=DEFAULT_URI,
+                overwrite=False,
+            )
+
+            # Create index
+            self.index = VectorStoreIndex.from_documents(
+                documents=documents,
+                embed_model=self.embeddings,
+                vector_store=self.vector_store,
+            )
+        else:
+            self.index = vectorstore
+
+    async def retrieve(self, query: str, k: int = 1) -> list[str]:
+        """
+        Retrieves similar documents based on the given query using LlamaIndex.
+
+        Args:
+            query (str): The query string to search for similar documents.
+            k (int, optional): The number of similar documents to retrieve.
+            Defaults to 1.
+
+        Returns:
+            list[str]: A list of document contents that are similar to the
+            query.
+
+        Raises:
+            ValueError: If the index is not initialized.
+        """
+        if self.index is None:
+            raise ValueError(
+                "Index not initialized.",
+            )
+
+        # Create query engine and query
+        query_engine = self.index.as_retriever(similarity_top_k=k)
+        response = query_engine.retrieve(query)
+
+        # Extract text from nodes
+        if len(response) > 0:
+            return [node.node.get_content() for node in response]
+        else:
+            return [""]

agentscope_runtime/engine/services/reme_personal_memory_service.py

@@ -0,0 +1,106 @@
+# -*- coding: utf-8 -*-
+import os
+from typing import Optional, Dict, Any, List
+
+from .memory_service import MemoryService
+from ..schemas.agent_schemas import Message
+
+
+class ReMePersonalMemoryService(MemoryService):
+    """
+    ReMe requires the following env variables to be set:
+    FLOW_EMBEDDING_API_KEY=sk-xxxx
+    FLOW_EMBEDDING_BASE_URL=https://xxxx/v1
+    FLOW_LLM_API_KEY=sk-xxxx
+    FLOW_LLM_BASE_URL=https://xxxx/v1
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        for key in [
+            "FLOW_EMBEDDING_API_KEY",
+            "FLOW_EMBEDDING_BASE_URL",
+            "FLOW_LLM_API_KEY",
+            "FLOW_LLM_BASE_URL",
+        ]:
+            if os.getenv(key) is None:
+                raise ValueError(f"{key} is not set")
+
+        from reme_ai.service.personal_memory_service import (
+            PersonalMemoryService,
+        )
+
+        self.service = PersonalMemoryService()
+
+    @staticmethod
+    def transform_message(message: Message) -> dict:
+        content_text = None
+
+        try:
+            if hasattr(message, "content") and isinstance(
+                message.content,
+                list,
+            ):
+                if len(message.content) > 0 and hasattr(
+                    message.content[0],
+                    "text",
+                ):
+                    content_text = message.content[0].text
+        except (AttributeError, IndexError):
+            # Log error or handle appropriately
+            pass
+
+        return {
+            "role": getattr(message, "role", None),
+            "content": content_text,
+        }
+
+    def transform_messages(self, messages: List[Message]) -> List[dict]:
+        return [self.transform_message(message) for message in messages]
+
+    async def start(self) -> None:
+        return await self.service.start()
+
+    async def stop(self) -> None:
+        return await self.service.stop()
+
+    async def health(self) -> bool:
+        return await self.service.health()
+
+    async def add_memory(
+        self,
+        user_id: str,
+        messages: list,
+        session_id: Optional[str] = None,
+    ) -> None:
+        return await self.service.add_memory(
+            user_id,
+            self.transform_messages(messages),
+            session_id,
+        )
+
+    async def search_memory(
+        self,
+        user_id: str,
+        messages: list,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> list:
+        return await self.service.search_memory(
+            user_id,
+            self.transform_messages(messages),
+            filters,
+        )
+
+    async def list_memory(
+        self,
+        user_id: str,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> list:
+        return await self.service.list_memory(user_id, filters)
+
+    async def delete_memory(
+        self,
+        user_id: str,
+        session_id: Optional[str] = None,
+    ) -> None:
+        return await self.service.delete_memory(user_id, session_id)

agentscope_runtime/engine/services/reme_task_memory_service.py

@@ -0,0 +1,11 @@
+# -*- coding: utf-8 -*-
+from .reme_personal_memory_service import ReMePersonalMemoryService
+
+
+class ReMeTaskMemoryService(ReMePersonalMemoryService):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        from reme_ai.service.task_memory_service import TaskMemoryService
+
+        self.service = TaskMemoryService()