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

agentscope_runtime/engine/agents/agentscope_agent/agent.py

@@ -85,6 +85,7 @@ class AgentScopeContextAdapter:
 
     @staticmethod
     def converter(message: Message):
+        # TODO: support more message type
         if message.role not in ["user", "system", "assistant"]:
             role_label = "user"
         else:

agentscope_runtime/engine/agents/agno_agent.py

@@ -53,6 +53,7 @@ class AgnoContextAdapter:
 
     @staticmethod
     def converter(message: Message):
+        # TODO: support more message type
         return dict(message)
 
     async def adapt_new_message(self):

agentscope_runtime/engine/agents/autogen_agent.py

@@ -0,0 +1,245 @@
+# -*- coding: utf-8 -*-
+from typing import Optional, Type
+
+from autogen_core.models import ChatCompletionClient
+from autogen_core.tools import FunctionTool
+from autogen_agentchat.agents import AssistantAgent
+from autogen_agentchat.messages import (
+    TextMessage,
+    ToolCallExecutionEvent,
+    ToolCallRequestEvent,
+    ModelClientStreamingChunkEvent,
+)
+
+from ..agents import Agent
+from ..schemas.context import Context
+from ..schemas.agent_schemas import (
+    Message,
+    TextContent,
+    DataContent,
+    FunctionCall,
+    FunctionCallOutput,
+    MessageType,
+    RunStatus,
+)
+
+
+class AutogenContextAdapter:
+    def __init__(self, context: Context, attr: dict):
+        self.context = context
+        self.attr = attr
+
+        # Adapted attribute
+        self.toolkit = None
+        self.model = None
+        self.memory = None
+        self.new_message = None
+
+    async def initialize(self):
+        self.model = await self.adapt_model()
+        self.memory = await self.adapt_memory()
+        self.new_message = await self.adapt_new_message()
+        self.toolkit = await self.adapt_tools()
+
+    async def adapt_memory(self):
+        messages = []
+
+        # Build context
+        for msg in self.context.session.messages[:-1]:  # Exclude the last one
+            messages.append(AutogenContextAdapter.converter(msg))
+
+        return messages
+
+    @staticmethod
+    def converter(message: Message):
+        # TODO: support more message type
+        return TextMessage.load(
+            {
+                "id": message.id,
+                "source": message.role,
+                "content": message.content[0].text if message.content else "",
+            },
+        )
+
+    async def adapt_new_message(self):
+        last_message = self.context.session.messages[-1]
+
+        return AutogenContextAdapter.converter(last_message)
+
+    async def adapt_model(self):
+        return self.attr["model"]
+
+    async def adapt_tools(self):
+        toolkit = self.attr["agent_config"].get("toolkit", [])
+        tools = self.attr["tools"]
+
+        # in case, tools is None and tools == []
+        if not tools:
+            return toolkit
+
+        if self.context.activate_tools:
+            # Only add activated tool
+            activated_tools = self.context.activate_tools
+        else:
+            from ...sandbox.tools.utils import setup_tools
+
+            activated_tools = setup_tools(
+                tools=self.attr["tools"],
+                environment_manager=self.context.environment_manager,
+                session_id=self.context.session.id,
+                user_id=self.context.session.user_id,
+                include_schemas=False,
+            )
+
+        for tool in activated_tools:
+            func = FunctionTool(
+                func=tool.make_function(),
+                description=tool.schema["function"]["description"],
+                name=tool.name,
+            )
+            toolkit.append(func)
+
+        return toolkit
+
+
+class AutogenAgent(Agent):
+    def __init__(
+        self,
+        name: str,
+        model: ChatCompletionClient,
+        tools=None,
+        agent_config=None,
+        agent_builder: Optional[Type[AssistantAgent]] = AssistantAgent,
+    ):
+        super().__init__(name=name, agent_config=agent_config)
+
+        assert isinstance(
+            model,
+            ChatCompletionClient,
+        ), "model must be a subclass of ChatCompletionClient in autogen"
+
+        # Set default agent_builder
+        if agent_builder is None:
+            agent_builder = AssistantAgent
+
+        assert issubclass(
+            agent_builder,
+            AssistantAgent,
+        ), "agent_builder must be a subclass of AssistantAgent in autogen"
+
+        # Replace name if not exists
+        self.agent_config["name"] = self.agent_config.get("name") or name
+
+        self._attr = {
+            "model": model,
+            "tools": tools,
+            "agent_config": self.agent_config,
+            "agent_builder": agent_builder,
+        }
+        self._agent = None
+        self.tools = tools
+
+    def copy(self) -> "AutogenAgent":
+        return AutogenAgent(**self._attr)
+
+    def build(self, as_context):
+        self._agent = self._attr["agent_builder"](
+            **self._attr["agent_config"],
+            model_client=as_context.model,
+            tools=as_context.toolkit,
+        )
+
+        return self._agent
+
+    async def run(self, context):
+        ag_context = AutogenContextAdapter(context=context, attr=self._attr)
+        await ag_context.initialize()
+
+        # We should always build a new agent since the state is manage outside
+        # the agent
+        self._agent = self.build(ag_context)
+
+        resp = self._agent.run_stream(
+            task=ag_context.memory + [ag_context.new_message],
+        )
+
+        text_message = Message(
+            type=MessageType.MESSAGE,
+            role="assistant",
+            status=RunStatus.InProgress,
+        )
+        yield text_message
+
+        text_delta_content = TextContent(delta=True)
+        is_text_delta = False
+        stream_mode = False
+        async for event in resp:
+            if getattr(event, "source", "user") == "user":
+                continue
+
+            if isinstance(event, TextMessage):
+                if stream_mode:
+                    continue
+                is_text_delta = True
+                text_delta_content.text = event.content
+                text_delta_content = text_message.add_delta_content(
+                    new_content=text_delta_content,
+                )
+                yield text_delta_content
+            elif isinstance(event, ModelClientStreamingChunkEvent):
+                stream_mode = True
+                is_text_delta = True
+                text_delta_content.text = event.content
+                text_delta_content = text_message.add_delta_content(
+                    new_content=text_delta_content,
+                )
+                yield text_delta_content
+            elif isinstance(event, ToolCallRequestEvent):
+                data = DataContent(
+                    data=FunctionCall(
+                        call_id=event.id,
+                        name=event.content[0].name,
+                        arguments=event.content[0].arguments,
+                    ).model_dump(),
+                )
+                message = Message(
+                    type=MessageType.PLUGIN_CALL,
+                    role="assistant",
+                    status=RunStatus.Completed,
+                    content=[data],
+                )
+                yield message
+            elif isinstance(event, ToolCallExecutionEvent):
+                data = DataContent(
+                    data=FunctionCallOutput(
+                        call_id=event.id,
+                        output=event.content[0].content,
+                    ).model_dump(),
+                )
+                message = Message(
+                    type=MessageType.PLUGIN_CALL_OUTPUT,
+                    role="assistant",
+                    status=RunStatus.Completed,
+                    content=[data],
+                )
+                yield message
+
+                # Add to message
+                is_text_delta = True
+                text_delta_content.text = event.content[0].content
+                text_delta_content = text_message.add_delta_content(
+                    new_content=text_delta_content,
+                )
+                yield text_delta_content
+
+        if is_text_delta:
+            yield text_message.content_completed(text_delta_content.index)
+            yield text_message.completed()
+
+    async def run_async(
+        self,
+        context,
+        **kwargs,
+    ):
+        async for event in self.run(context):
+            yield event

agentscope_runtime/engine/schemas/agent_schemas.py

@@ -210,7 +210,7 @@ class Content(Event):
     delta: Optional[bool] = False
     """Whether this content is a delta."""
 
-    msg_id: str = None
+    msg_id: Optional[str] = None
     """message unique id"""
 
 

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,18 @@ 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)],
+            )
+            session.messages.append(message)
+
 
 class ContextManager(ServiceManager):
     """
@@ -53,10 +73,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 +90,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 +101,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 +144,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/memory_service.py

@@ -58,8 +58,8 @@ class MemoryService(ServiceWithLifecycleManager):
         Args:
             user_id: The user id.
             messages: The user query or the query with history messages,
-            both in the format of list of messages.  If messages is a list,
-            the search will be based on the content of the last message.
+                both in the format of list of messages.  If messages is a list,
+                the search will be based on the content of the last message.
             filters: The filters used to search memory
         """
 

agentscope_runtime/engine/services/rag_service.py

@@ -0,0 +1,101 @@
+# -*- coding: utf-8 -*-
+from typing import Optional
+
+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]:
+        raise NotImplementedError
+
+
+DEFAULT_URI = "milvus_demo.db"
+
+
+class LangChainRAGService(RAGService):
+    """
+    RAG Service using LangChain
+    """
+
+    def __init__(
+        self,
+        uri: Optional[str] = None,
+        docs: Optional[list[str]] = 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)
+        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"},
+        )
+
+    async def retrieve(self, query: str, k: int = 1) -> list[str]:
+        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

agentscope_runtime/engine/services/redis_memory_service.py

@@ -0,0 +1,187 @@
+# -*- coding: utf-8 -*-
+from typing import Optional, Dict, Any
+import json
+import redis.asyncio as aioredis
+
+
+from .memory_service import MemoryService
+from ..schemas.agent_schemas import Message, MessageType
+
+
+class RedisMemoryService(MemoryService):
+    """
+    A Redis-based implementation of the memory service.
+    """
+
+    def __init__(
+        self,
+        redis_url: str = "redis://localhost:6379/0",
+        redis_client: Optional[aioredis.Redis] = None,
+    ):
+        self._redis_url = redis_url
+        self._redis = redis_client
+        self._DEFAULT_SESSION_ID = "default"
+
+    async def start(self) -> None:
+        """Starts the Redis connection."""
+        if self._redis is None:
+            self._redis = aioredis.from_url(
+                self._redis_url,
+                decode_responses=True,
+            )
+
+    async def stop(self) -> None:
+        """Closes the Redis connection."""
+        if self._redis:
+            await self._redis.close()
+            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 _user_key(self, user_id):
+        # Each user is a Redis hash
+        return f"user_memory:{user_id}"
+
+    def _serialize(self, messages):
+        return json.dumps([msg.dict() for msg in messages])
+
+    def _deserialize(self, messages_json):
+        if not messages_json:
+            return []
+        return [Message.parse_obj(m) for m in json.loads(messages_json)]
+
+    async def add_memory(
+        self,
+        user_id: str,
+        messages: list,
+        session_id: Optional[str] = None,
+    ) -> None:
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
+        key = self._user_key(user_id)
+        field = session_id if session_id else self._DEFAULT_SESSION_ID
+
+        existing_json = await self._redis.hget(key, field)
+        existing_msgs = self._deserialize(existing_json)
+        all_msgs = existing_msgs + messages
+        await self._redis.hset(key, field, self._serialize(all_msgs))
+
+    async def search_memory(
+        self,
+        user_id: str,
+        messages: list,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> list:
+        key = self._user_key(user_id)
+        if (
+            not messages
+            or not isinstance(messages, list)
+            or len(messages) == 0
+        ):
+            return []
+
+        message = messages[-1]
+        query = await self.get_query_text(message)
+        if not query:
+            return []
+
+        keywords = set(query.lower().split())
+
+        all_msgs = []
+        hash_keys = await self._redis.hkeys(key)
+        for session_id in hash_keys:
+            msgs_json = await self._redis.hget(key, session_id)
+            msgs = self._deserialize(msgs_json)
+            all_msgs.extend(msgs)
+
+        matched_messages = []
+        for msg in all_msgs:
+            candidate_content = await self.get_query_text(msg)
+            if candidate_content:
+                msg_content_lower = candidate_content.lower()
+                if any(keyword in msg_content_lower for keyword in keywords):
+                    matched_messages.append(msg)
+
+        if (
+            filters
+            and "top_k" in filters
+            and isinstance(filters["top_k"], int)
+        ):
+            return matched_messages[-filters["top_k"] :]
+
+        return matched_messages
+
+    async def get_query_text(self, message: Message) -> str:
+        if message:
+            if message.type == MessageType.MESSAGE:
+                for content in message.content:
+                    if content.type == "text":
+                        return content.text
+        return ""
+
+    async def list_memory(
+        self,
+        user_id: str,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> list:
+        key = self._user_key(user_id)
+        all_msgs = []
+        hash_keys = await self._redis.hkeys(key)
+        for session_id in sorted(hash_keys):
+            msgs_json = await self._redis.hget(key, session_id)
+            msgs = self._deserialize(msgs_json)
+            all_msgs.extend(msgs)
+
+        page_num = filters.get("page_num", 1) if filters else 1
+        page_size = filters.get("page_size", 10) if filters else 10
+
+        start_index = (page_num - 1) * page_size
+        end_index = start_index + page_size
+
+        return all_msgs[start_index:end_index]
+
+    async def delete_memory(
+        self,
+        user_id: str,
+        session_id: Optional[str] = None,
+    ) -> None:
+        key = self._user_key(user_id)
+        if session_id:
+            await self._redis.hdel(key, session_id)
+        else:
+            await self._redis.delete(key)
+
+    async def clear_all_memory(self) -> None:
+        """
+        Clears all memory data from Redis.
+        This method removes all user memory keys from the Redis database.
+        """
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
+
+        keys = await self._redis.keys(self._user_key("*"))
+        if keys:
+            await self._redis.delete(*keys)
+
+    async def delete_user_memory(self, user_id: str) -> None:
+        """
+        Deletes all memory data for a specific user.
+
+        Args:
+            user_id (str): The ID of the user whose memory data should be
+            deleted
+        """
+        if not self._redis:
+            raise RuntimeError("Redis connection is not available")
+
+        key = self._user_key(user_id)
+        await self._redis.delete(key)