# RAG Module Architecture in AgentScope: A Deep Dive into Retrieval-Augmented Generation

> Explore the RAG module architecture in AgentScope. Learn how it transforms data into vector indexes for efficient retrieval during agent conversations.

- Repository: [AgentScope-AI/agentscope](https://github.com/agentscope-ai/agentscope)
- Tags: deep-dive
- Published: 2026-03-09

---

**The RAG module in AgentScope is a modular, plug-in-friendly pipeline that converts raw data into vector-indexed documents, stores them in configurable vector databases, and enables agents to retrieve relevant context on-the-fly during conversations.**

AgentScope's retrieval-augmented generation subsystem is designed as a decoupled pipeline that separates data ingestion, embedding, storage, and retrieval. According to the agentscope-ai/agentscope source code, this architecture allows developers to swap individual components—such as readers, embedding models, or vector stores—without modifying the core retrieval logic or agent behavior.

## Core Components of the RAG Module

The RAG architecture consists of six primary abstractions that work together to move data from raw files to agent context.

### Readers and Document Ingestion

**Readers** handle the initial data ingestion. Each reader inherits from `ReaderBase` defined in [`src/agentscope/rag/_reader/_reader_base.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_reader/_reader_base.py) and implements two key methods: `__call__` (async) for processing input and `get_doc_id` for generating unique identifiers. Concrete implementations like `TextReader` (located in [`src/agentscope/rag/_reader/_text_reader.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_reader/_text_reader.py)) handle specific formats including PDF, PPT, images, Excel, and plain text, converting them into standardized `Document` objects.

### Document Structure

The **Document** class in [`src/agentscope/rag/_document.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_document.py) serves as the atomic unit of knowledge. Each document contains raw text content, a `DocMetadata` object for tracking source information, and an optional embedding vector. This structure ensures that chunks maintain provenance metadata as they move through the pipeline.

### Knowledge Base Abstraction

The **KnowledgeBase** abstract class in [`src/agentscope/rag/_knowledge_base.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_knowledge_base.py) defines the high-level API for document management. It specifies methods for adding documents (which triggers embedding) and retrieving relevant chunks via queries. The default implementation, **SimpleKnowledge** in [`src/agentscope/rag/_simple_knowledge.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_simple_knowledge.py), orchestrates the embedding model and vector store to handle the actual persistence and search operations.

### Vector Database Store

**VDBStoreBase** in [`src/agentscope/rag/_store/_store_base.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_store/_store_base.py) provides a thin, consistent wrapper around various vector databases including MilvusLite, Qdrant, MongoDB, OceanBase, and AlibabaCloud MySQL. The interface requires implementing `add`, `search`, `delete`, and `get_client` methods. Concrete stores like `MilvusLiteStore` handle database-specific connection logic while exposing a unified API to the knowledge base.

### Embedding Models

Any model implementing **EmbeddingModelBase** can generate dense vectors for both documents and queries. The system supports OpenAI, Gemini, Ollama, and custom embedding providers through this interface, allowing the same knowledge base to work with different vectorization strategies.

### Agent Integration

The **ReActAgent** in [`src/agentscope/agent/_react_agent.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/agent/_react_agent.py) integrates retrieval directly into the conversation loop via the `_retrieve_from_knowledge` method. This method automatically extracts the user query, optionally rewrites it using the agent's LLM for better retrieval, queries registered knowledge bases, and injects the results into the agent's short-term memory as `<retrieved_knowledge>` blocks.

## Data Flow Through the Pipeline

The RAG module processes information through four distinct stages:

1. **Ingestion**: A reader's `__call__` method processes raw input (files, URLs, or text) and returns a list of `Document` objects with populated `DocMetadata`.

2. **Embedding and Storage**: `SimpleKnowledge.add_documents(docs)` calls the configured `EmbeddingModel` to generate vectors, then persists the documents through the `VDBStore` implementation's `add` method.

3. **Retrieval**: When `ReActAgent._retrieve_from_knowledge` triggers, it optionally rewrites the query using `_QueryRewriteModel`, then calls `KnowledgeBase.retrieve(query)` which delegates to the store's `search` method to return ranked `Document` objects.

4. **Agent Reasoning**: Retrieved documents are wrapped as `TextBlock` objects inside a `ToolResponse` and injected into the prompt. The agent generates responses grounded in this external context via the standard `reply` flow.

## Implementation Examples

### Building a Knowledge Base from Plain Text

The following example demonstrates creating a knowledge base using the text reader, OpenAI embeddings, and MilvusLite storage:

```python
from agentscope.rag import TextReader, SimpleKnowledge
from agentscope.embedding import OpenAIEmbeddingModel
from agentscope.rag import MilvusLiteStore

# Reader splits text into 512-character chunks

reader = TextReader(chunk_size=512, split_by="sentence")
documents = await reader("docs/intro.md")

# Initialize embedding model and vector store

embedder = OpenAIEmbeddingModel(api_key="YOUR_KEY")
store = MilvusLiteStore(collection_name="my_collection")

# Create knowledge base and add documents

knowledge = SimpleKnowledge(embedding_store=store, embedding_model=embedder)
await knowledge.add_documents(documents)

```

### Querying from a ReActAgent

This configuration enables automatic retrieval during agent conversations:

```python
from agentscope.agent import ReActAgent
from agentscope.model import OpenAIChatModel
from agentscope.formatter import OpenAIFmt

model = OpenAIChatModel(api_key="YOUR_KEY")
formatter = OpenAIFmt()

agent = ReActAgent(
    name="assistant",
    sys_prompt="You are a helpful assistant.",
    model=model,
    formatter=formatter,
    knowledge=knowledge,
    enable_rewrite_query=True,
)

reply = await agent.reply("What are the three main contributions of the RAG module?")

```

When `enable_rewrite_query` is set to `True`, the agent first uses its LLM to optimize the user query before retrieval, then automatically injects the relevant documents into the reasoning context.

### Adding a Custom Vector Store

To integrate a proprietary vector database, subclass `VDBStoreBase` and implement the required interface:

```python
from agentscope.rag import VDBStoreBase

class MyVectorStore(VDBStoreBase):
    async def add(self, documents, **kwargs):
        # Implementation for adding vectors

        pass
    
    async def search(self, query_embedding, limit, score_threshold=None, **kwargs):
        # Implementation for similarity search

        pass
    
    async def delete(self, *args, **kwargs):
        # Implementation for deletion

        pass

# Use the custom store

my_store = MyVectorStore()
knowledge = SimpleKnowledge(embedding_store=my_store, embedding_model=embedder)

```

## Summary

- **The RAG module in AgentScope** uses a pipeline architecture with clear separation between ingestion, embedding, storage, and retrieval.
- **Readers** in `src/agentscope/rag/_reader/` handle file parsing and chunking, producing `Document` objects defined in [`src/agentscope/rag/_document.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_document.py).
- **SimpleKnowledge** in [`src/agentscope/rag/_simple_knowledge.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_simple_knowledge.py) provides the default orchestration between embedding models and vector stores.
- **Vector stores** implement `VDBStoreBase` from [`src/agentscope/rag/_store/_store_base.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_store/_store_base.py), supporting multiple backends including MilvusLite and Qdrant.
- **ReActAgent** automatically retrieves knowledge via `_retrieve_from_knowledge` in [`src/agentscope/agent/_react_agent.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/agent/_react_agent.py), optionally rewriting queries and injecting results into the conversation context.
- The entire system is extensible through subclassing, allowing custom readers, stores, and embedding models without modifying core agent logic.

## Frequently Asked Questions

### How does AgentScope handle different file formats in the RAG pipeline?

AgentScope delegates file parsing to specialized **Reader** classes. The `ReaderBase` abstract class in [`src/agentscope/rag/_reader/_reader_base.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_reader/_reader_base.py) defines the interface, while concrete implementations like `TextReader`, `PdfReader`, and `ImageReader` handle specific formats. Each reader converts raw inputs into standardized `Document` objects, allowing the rest of the pipeline to remain format-agnostic.

### What is the difference between KnowledgeBase and SimpleKnowledge in AgentScope?

`KnowledgeBase` in [`src/agentscope/rag/_knowledge_base.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_knowledge_base.py) is the abstract interface that defines the contract for adding and retrieving documents. `SimpleKnowledge` in [`src/agentscope/rag/_simple_knowledge.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_simple_knowledge.py) is the concrete default implementation that combines an `EmbeddingModel` with a `VDBStore` to handle the actual vectorization and persistence. Developers can subclass `KnowledgeBase` to create custom retrieval logic while keeping the same agent-facing API.

### Can I use a custom vector database with the AgentScope RAG module?

Yes. The `VDBStoreBase` class in [`src/agentscope/rag/_store/_store_base.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/rag/_store/_store_base.py) provides an abstraction layer requiring only four methods: `add`, `search`, `delete`, and `get_client`. By subclassing this base and implementing these methods for your specific database (whether proprietary or specialized), you can pass your custom store directly to `SimpleKnowledge` without changing any reader or agent code.

### How does query rewriting improve retrieval in AgentScope agents?

When `enable_rewrite_query` is set to `True` in `ReActAgent`, the system uses the agent's LLM to transform the raw user query into an embedding-friendly format before searching the vector store. This happens in `_retrieve_from_knowledge` in [`src/agentscope/agent/_react_agent.py`](https://github.com/agentscope-ai/agentscope/blob/main/src/agentscope/agent/_react_agent.py), where the `_QueryRewriteModel` generates alternative phrasings or expansions that better match the embedded document chunks, significantly improving recall for complex or ambiguous questions.