add rag service

Author: XGenerationLab

src/agentscope_runtime/engine/services/context_manager.py

@@ -4,12 +4,19 @@
 
 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 @@ async def compose(
         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 @@ async def compose(
             )
             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 @@ def __init__(
         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,7 @@ def _register_default_services(self):
 
         self.register_service("session", self._session_history_service)
         self.register_service("memory", self._memory_service)
+        self.register_service("rag", self._rag_service)
 
     async def compose_context(
         self,
@@ -77,6 +100,7 @@ async def compose_context(
         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 +143,12 @@ async def append(self, session: Session, event_output: List[Message]):
 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:

src/agentscope_runtime/engine/services/rag_service.py

@@ -4,13 +4,34 @@
 from langchain_milvus import Milvus
 
 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"
 
@@ -31,7 +52,9 @@ def __init__(self, uri=None, docs=None):
             self.uri = DEFAULT_URI
             self.from_docs(docs)
         else:
-            raise ValueError("Either uri or docs must be provided.")
+            docs = []
+            self.uri = DEFAULT_URI
+            self.from_docs(docs)
 
     def from_docs(self, docs=None):
         if docs is None:
@@ -53,17 +76,16 @@ def from_db(self):
             index_params={"index_type": "FLAT", "metric_type": "L2"},
         )
 
-    async def retrieve(self, query: str, k: int = 1) -> list:
+    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.",
             )
-        return self.vectorstore.similarity_search(query, k=k)
+        docs = self.vectorstore.similarity_search(query, k=k)
+        return [doc.page_content for doc in docs]
 
     async def start(self) -> None:
         """Starts the service."""
-        self.embeddings = DashScopeEmbeddings()
-        self.vectorstore = None
 
     async def stop(self) -> None:
         """Stops the service."""

tests/unit/test_rag_service.py

@@ -4,6 +4,17 @@
 import pytest
 from dotenv import load_dotenv
 
+from agentscope_runtime.engine import Runner
+from agentscope_runtime.engine.agents.llm_agent import LLMAgent
+from agentscope_runtime.engine.llms import QwenLLM
+from agentscope_runtime.engine.schemas.agent_schemas import (
+    MessageType,
+    AgentRequest,
+    RunStatus,
+)
+from agentscope_runtime.engine.services.context_manager import (
+    create_context_manager,
+)
 from agentscope_runtime.engine.services.rag_service import LangChainRAGService
 
 if os.path.exists("../../.env"):
@@ -46,7 +57,7 @@ async def test_from_docs():
         "What is self-reflection of an AI Agent?",
     )
     assert len(ret_docs) == 1
-    assert ret_docs[0].page_content.startswith("Self-Reflection")
+    assert ret_docs[0].startswith("Self-Reflection")
 
 
 @pytest.mark.asyncio
@@ -56,4 +67,56 @@ async def test_from_db():
         "What is self-reflection of an AI Agent?",
     )
     assert len(ret_docs) == 1
-    assert ret_docs[0].page_content.startswith("Self-Reflection")
+    assert ret_docs[0].startswith("Self-Reflection")
+
+
+@pytest.mark.asyncio
+async def test_rag():
+    rag_service = LangChainRAGService(uri="./assets/milvus_demo.db")
+    USER_ID = "user2"
+    SESSION_ID = "session1"
+    query = "What is self-reflection of an AI Agent?"
+
+    llm_agent = LLMAgent(
+        model=QwenLLM(),
+        name="llm_agent",
+        description="A simple LLM agent",
+    )
+
+    async with create_context_manager(
+        rag_service=rag_service,
+    ) as context_manager:
+        runner = Runner(
+            agent=llm_agent,
+            context_manager=context_manager,
+            environment_manager=None,
+        )
+
+        all_result = ""
+        # print("\n")
+        request = AgentRequest(
+            input=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": query,
+                        },
+                    ],
+                },
+            ],
+            session_id=SESSION_ID,
+        )
+
+        async for message in runner.stream_query(
+            user_id=USER_ID,
+            request=request,
+        ):
+            if (
+                message.object == "message"
+                and MessageType.MESSAGE == message.type
+                and RunStatus.Completed == message.status
+            ):
+                all_result = message.content[0].text
+        print(all_result)