agentscope-runtime 0.1.1__py3-none-any.whl → 0.1.3__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

@@ -4,12 +4,19 @@ from typing import List
 
 from .manager import ServiceManager
 from .memory_service import MemoryService, InMemoryMemoryService
+from .rag_service import RAGService
 from .session_history_service import (
     SessionHistoryService,
     Session,
     InMemorySessionHistoryService,
 )
-from ..schemas.agent_schemas import Message
+from ..schemas.agent_schemas import (
+    Message,
+    MessageType,
+    Role,
+    TextContent,
+    ContentType,
+)
 
 
 class ContextComposer:
@@ -19,6 +26,7 @@ class ContextComposer:
         session: Session,  # session
         memory_service: MemoryService = None,
         session_history_service: SessionHistoryService = None,
+        rag_service: RAGService = None,
     ):
         # session
         if session_history_service:
@@ -42,6 +50,24 @@ class ContextComposer:
             )
             session.messages = memories + session.messages
 
+        # rag
+        if rag_service:
+            query = await rag_service.get_query_text(request_input[-1])
+            docs = await rag_service.retrieve(query=query, k=5)
+            cooked_doc = "\n".join(docs)
+            message = Message(
+                type=MessageType.MESSAGE,
+                role=Role.SYSTEM,
+                content=[TextContent(type=ContentType.TEXT, text=cooked_doc)],
+            )
+            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):
     """
@@ -53,10 +79,12 @@ class ContextManager(ServiceManager):
         context_composer_cls=ContextComposer,
         session_history_service: SessionHistoryService = None,
         memory_service: MemoryService = None,
+        rag_service: RAGService = None,
     ):
         self._context_composer_cls = context_composer_cls
         self._session_history_service = session_history_service
         self._memory_service = memory_service
+        self._rag_service = rag_service
         super().__init__()
 
     def _register_default_services(self):
@@ -68,6 +96,8 @@ class ContextManager(ServiceManager):
 
         self.register_service("session", self._session_history_service)
         self.register_service("memory", self._memory_service)
+        if self._rag_service:
+            self.register_service("rag", self._rag_service)
 
     async def compose_context(
         self,
@@ -77,6 +107,7 @@ class ContextManager(ServiceManager):
         await self._context_composer_cls.compose(
             memory_service=self._memory_service,
             session_history_service=self._session_history_service,
+            rag_service=self._rag_service,
             session=session,
             request_input=request_input,
         )
@@ -119,10 +150,12 @@ class ContextManager(ServiceManager):
 async def create_context_manager(
     memory_service: MemoryService = None,
     session_history_service: SessionHistoryService = None,
+    rag_service: RAGService = None,
 ):
     manager = ContextManager(
         memory_service=memory_service,
         session_history_service=session_history_service,
+        rag_service=rag_service,
     )
 
     async with manager:

agentscope_runtime/engine/services/rag_service.py

@@ -0,0 +1,195 @@
+# -*- coding: utf-8 -*-
+from .base import ServiceWithLifecycleManager
+from ..schemas.agent_schemas import Message, MessageType
+
+
+class RAGService(ServiceWithLifecycleManager):
+    """
+    RAG Service
+    """
+
+    async def get_query_text(self, message: Message) -> str:
+        """
+        Gets the query text from the messages.
+
+        Args:
+            message: A list of messages.
+
+        Returns:
+            The query text.
+        """
+        if message:
+            if message.type == MessageType.MESSAGE:
+                for content in message.content:
+                    if content.type == "text":
+                        return content.text
+        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"
+
+
+class LangChainRAGService(RAGService):
+    """
+    RAG Service using LangChain
+    """
+
+    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 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.",
+            )
+        docs = self.vectorstore.similarity_search(query, k=k)
+        return [doc.page_content for doc in docs]
+
+    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
+
+
+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 [""]