Add RAG module into AgentScope

examples/functionality/rag/README.md

@@ -0,0 +1,11 @@
+# RAG in AgentScope
+
+This example includes two scripts to demonstrate how to use Retrieval-Augmented Generation (RAG) in AgentScope:
+
+- the basic usage of RAG module in AgentScope in ``basic_usage.py``
+- a simple agentic use case of RAG in ``agentic_usage.py``
+
+We also provide a solution to integrate RAG into the `ReActAgent` class by retrieving relevant documents
+at the beginning of each reply in a ``pre_reply`` hook.
+Because it's too similar with the long-term memory example, we are considering if integrating the RAG module
+into the `ReActAgent` class by supporting a `knowledge` argument in the constructor.

examples/functionality/rag/agentic_usage.py

@@ -0,0 +1,146 @@
+# -*- coding: utf-8 -*-
+"""The agentic usage example for RAG in AgentScope, where the agent is
+equipped with RAG tools to answer questions based on a knowledge base.
+
+The example is more challenging for the agent, requiring the agent to
+adjust the retrieval parameters to get relevant results.
+"""
+import asyncio
+import os
+
+from agentscope.agent import ReActAgent, UserAgent
+from agentscope.embedding import DashScopeTextEmbedding
+from agentscope.formatter import DashScopeChatFormatter
+from agentscope.message import Msg, TextBlock
+from agentscope.model import DashScopeChatModel
+from agentscope.rag import SimpleKnowledge, QdrantStore, TextReader
+from agentscope.tool import Toolkit, ToolResponse
+
+# Create a knowledge base instance
+knowledge = SimpleKnowledge(
+    embedding_store=QdrantStore(
+        location=":memory:",
+        collection_name="test_collection",
+        dimensions=1024,  # The dimension of the embedding vectors
+    ),
+    embedding_model=DashScopeTextEmbedding(
+        api_key=os.environ["DASHSCOPE_API_KEY"],
+        model_name="text-embedding-v4",
+    ),
+)
+
+
+# An agent-oriented tool function to retrieve knowledge from the knowledge base
+# We expose the `limit` and `score_threshold` parameters to the agent
+# for more flexible control
+async def retrieve_knowledge(
+    query: str,
+    limit: int = 3,
+    score_threshold: float | None = 0.8,
+) -> ToolResponse:
+    """Retrieve relevant documents from the knowledge base, which is relevant
+    to John Doe's profile. Note the `query` parameter is
+    very important for the retrieval quality, and you can try many different
+    queries to get the best results. Adjust the `limit` and `score_threshold`
+    parameters to get more or fewer results.
+
+    Args:
+        query (`str`):
+            The query string, which should be specific and concise. For
+            example, you should provide the specific name instead of
+            "you", "my", "he", "she", etc.
+        limit (`int`, defaults to 3):
+            The number of relevant documents to retrieve.
+        score_threshold (`float`, defaults to 0.8):
+            A threshold in [0, 1] and only the relevance score above this
+            threshold will be returned. Reduce this value to get more results.
+    """
+    # Retrieve relevant documents based on the given query
+    docs = await knowledge.retrieve(
+        query=query,
+        limit=limit,
+        score_threshold=score_threshold,
+    )
+    if len(docs):
+        return ToolResponse(
+            content=[
+                TextBlock(
+                    type="text",
+                    text=f"Score: {_.score}, "
+                    f"Content: {_.metadata.content['text']}",
+                )
+                for _ in docs
+            ],
+        )
+    return ToolResponse(
+        content=[
+            TextBlock(
+                type="text",
+                text="No relevant documents found. TRY to reduce the "
+                "`score_threshold` parameter to get "
+                "more results.",
+            ),
+        ],
+    )
+
+
+async def main() -> None:
+    """The main entry of the agent usage example for RAG in AgentScope."""
+
+    # Store some things into the knowledge base for demonstration
+    # In practice, the VDB store would be pre-filled with relevant data
+    reader = TextReader(chunk_size=1024, split_by="sentence")
+    documents = await reader(
+        text=(
+            # Fake personal profile for demonstration
+            "I'm John Doe, 28 years old. My best friend is James "
+            "Smith. I live in San Francisco. I work at OpenAI as a "
+            "software engineer. I love hiking and photography. "
+            "My father is Michael Doe, a doctor. I'm very proud of him. "
+            "My mother is Sarah Doe, a teacher. She is very kind and "
+            "always helps me with my studies.\n"
+            "I'm now a PhD student at Stanford University, majoring in "
+            "Computer Science. My advisor is Prof. Jane Williams, who is "
+            "a leading expert in artificial intelligence. I have published "
+            "several papers in top conferences, such as NeurIPS and ICML. "
+        ),
+    )
+    await knowledge.add_documents(documents)
+
+    # Create a toolkit and register the RAG tool function
+    toolkit = Toolkit()
+    toolkit.register_tool_function(retrieve_knowledge)
+
+    # Create an agent and a user
+    agent = ReActAgent(
+        name="Friday",
+        sys_prompt=(
+            "You're a helpful assistant named Friday. "
+            "You're equipped with a 'retrieve_knowledge' tool to help you "
+            "know about the user named John Doe. "
+            "NOTE to adjust the `score_threshold` parameters when you cannot "
+            "get relevant results. "
+        ),
+        toolkit=toolkit,
+        model=DashScopeChatModel(
+            api_key=os.environ["DASHSCOPE_API_KEY"],
+            model_name="qwen3-max-preview",
+        ),
+        formatter=DashScopeChatFormatter(),
+    )
+    user = UserAgent(name="User")
+
+    # A simple conversation loop beginning with a preset question
+    msg = Msg(
+        "user",
+        "I'm John Doe. Do you know my father?",
+        "user",
+    )
+    while True:
+        msg = await agent(msg)
+        msg = await user(msg)
+        if msg.get_text_content() == "exit":
+            break
+
+
+asyncio.run(main())

examples/functionality/rag/basic_usage.py

@@ -0,0 +1,79 @@
+# -*- coding: utf-8 -*-
+"""The main entry point of the RAG example."""
+import asyncio
+import os
+
+from agentscope.embedding import DashScopeTextEmbedding
+from agentscope.rag import (
+    TextReader,
+    PDFReader,
+    QdrantStore,
+    SimpleKnowledge,
+)
+
+
+async def main() -> None:
+    """The main entry point of the RAG example."""
+
+    # Create readers with chunking arguments
+    reader = TextReader(chunk_size=1024)
+    pdf_reader = PDFReader(chunk_size=1024, split_by="sentence")
+
+    # Read documents
+    documents = await reader(
+        text="I'm Tony Stank, my password is 123456. My best friend is James "
+        "Rhodes.",
+    )
+
+    # Read a sample PDF file
+    pdf_path = os.path.join(
+        os.path.abspath(os.path.dirname(__file__)),
+        "example.pdf",
+    )
+    pdf_documents = await pdf_reader(pdf_path=pdf_path)
+
+    # Create a knowledge base with Qdrant as the embedding store and
+    # DashScope as the embedding model
+    knowledge = SimpleKnowledge(
+        embedding_store=QdrantStore(
+            location=":memory:",
+            collection_name="test_collection",
+            dimensions=1024,  # The dimension of the embedding vectors
+        ),
+        embedding_model=DashScopeTextEmbedding(
+            api_key=os.environ["DASHSCOPE_API_KEY"],
+            model_name="text-embedding-v4",
+        ),
+    )
+
+    # Insert documents into the knowledge base
+    await knowledge.add_documents(documents + pdf_documents)
+
+    # Retrieve relevant documents based on a given query
+    docs = await knowledge.retrieve(
+        query="What is Tony Stank's password?",
+        limit=3,
+        score_threshold=0.7,
+    )
+    print("Q1: What is Tony Stank's password?")
+    for doc in docs:
+        print(
+            f"Document ID: {doc.id}, Score: {doc.score}, "
+            f"Content: {doc.metadata.content['text']}",
+        )
+
+    # Retrieve documents from the PDF file based on a query
+    docs = await knowledge.retrieve(
+        query="climate change",
+        limit=3,
+        score_threshold=0.2,
+    )
+    print("\n\nQ2: climate change")
+    for doc in docs:
+        print(
+            f"Document ID: {doc.id}, Score: {doc.score}, "
+            f"Content: {repr(doc.metadata.content['text'])}",
+        )
+
+
+asyncio.run(main())

examples/react_agent/main.py

@@ -17,6 +17,9 @@
 
 async def main() -> None:
     """The main entry point for the ReAct agent example."""
+    # import agentscope
+    #
+    # agentscope.init(studio_url="http://localhost:3000")
     toolkit = Toolkit()
     toolkit.register_tool_function(execute_shell_command)
     toolkit.register_tool_function(execute_python_code)
@@ -27,7 +30,7 @@ async def main() -> None:
         sys_prompt="You are a helpful assistant named Friday.",
         model=DashScopeChatModel(
             api_key=os.environ.get("DASHSCOPE_API_KEY"),
-            model_name="qwen-max",
+            model_name="qwen-vl-max",
             enable_thinking=False,
             stream=True,
         ),

setup.py

@@ -44,6 +44,7 @@
     "ray",
     "mem0ai",
     "packaging",
+    "pypdf",
 ]
 
 dev_requires = [

src/agentscope/agent/_react_agent.py

@@ -13,6 +13,7 @@
 from ..memory import MemoryBase, LongTermMemoryBase, InMemoryMemory
 from ..message import Msg, ToolUseBlock, ToolResultBlock, TextBlock
 from ..model import ChatModelBase
+from ..rag import KnowledgeBase
 from ..tool import Toolkit, ToolResponse
 from ..tracing import trace_reply
 
@@ -80,6 +81,7 @@ def __init__(
         ] = "both",
         enable_meta_tool: bool = False,
         parallel_tool_calls: bool = False,
+        knowledge: KnowledgeBase | list[KnowledgeBase] | None = None,
         max_iters: int = 10,
     ) -> None:
         """Initialize the ReAct agent
@@ -121,6 +123,9 @@ def __init__(
             parallel_tool_calls (`bool`, defaults to `False`):
                 When LLM generates multiple tool calls, whether to execute
                 them in parallel.
+            knowledge (`KnowledgeBase | list[KnowledgeBase] | None`, optional):
+                The knowledge object(s) used by the agent to retrieve
+                relevant documents at the beginning of each reply.
             max_iters (`int`, defaults to `10`):
                 The maximum number of iterations of the reasoning-acting loops.
         """
@@ -173,6 +178,12 @@ def __init__(
             )
 
         self.parallel_tool_calls = parallel_tool_calls
+
+        # The knowledge base(s) used by the agent
+        if isinstance(knowledge, KnowledgeBase):
+            knowledge = [knowledge]
+        self.knowledge: list[KnowledgeBase] = knowledge or []
+
         self.max_iters = max_iters
 
         # Variables to record the intermediate state

src/agentscope/embedding/__init__.py

@@ -5,6 +5,7 @@
 from ._embedding_usage import EmbeddingUsage
 from ._embedding_response import EmbeddingResponse
 from ._dashscope_embedding import DashScopeTextEmbedding
+from ._dashscope_multimodal_embedding import DashScopeMultiModalEmbedding
 from ._openai_embedding import OpenAITextEmbedding
 from ._gemini_embedding import GeminiTextEmbedding
 from ._ollama_embedding import OllamaTextEmbedding
@@ -17,6 +18,7 @@
     "EmbeddingUsage",
     "EmbeddingResponse",
     "DashScopeTextEmbedding",
+    "DashScopeMultiModalEmbedding",
     "OpenAITextEmbedding",
     "GeminiTextEmbedding",
     "OllamaTextEmbedding",