RAG Implementation Patterns for Agent Knowledge Base Integration: 5 Modular Architectures

The bojieli/ai-agent-book repository demonstrates a modular retrieval-augmented generation (RAG) architecture that unifies five distinct backends—ranging from offline BM25 to graph-based retrieval—behind a single LLM tool interface.

This article explores concrete RAG implementation patterns for agent knowledge base integration found in the open-source codebase. The architecture enables LLM-based agents to retrieve external knowledge through a unified tool contract, regardless of whether the data lives in a local Python index, a self-hosted API, or a cloud vector store.

The Unified Tool Contract Pattern

At the heart of the system lies a tool abstraction layer defined in chapter3/agentic-rag/tools.py. Rather than hardcoding retrieval logic into the agent, the codebase wraps each backend behind two standard functions: knowledge_base_search and get_document.

The get_tool_definitions() function (lines 75–94) exposes these capabilities to the LLM as JSON schema. When the agent invokes knowledge_base_search, the KnowledgeBaseTools class dispatches the query to the appropriate backend implementation based on runtime configuration. This pattern ensures that the LLM sees a consistent interface while the underlying retrieval mechanism can vary from BM25 to knowledge graphs.

Five RAG Implementation Patterns

The repository provides five concrete implementations, each optimized for different deployment scenarios and data structures.

Offline BM25 (Pure Python)

Best for: Local development, air-gapped environments, and small-to-medium corpora.

The offline pattern requires no external services. In KnowledgeBaseTools._search_offline, the system lazily initializes an OfflineRetriever (referenced in offline_retriever.py) that builds a BM25 index over a local text corpus. When the agent queries the knowledge base, the implementation ranks documents using sparse retrieval and returns a list of SearchResult objects containing chunk IDs, scores, and text snippets.

This pattern stores document metadata in a local JSON file specified by document_store_path and performs all computation in-memory.

Local Hybrid API (Self-Hosted Dense + Sparse)

Best for: Production deployments requiring full data sovereignty.

The local hybrid pattern connects to a self-hosted retrieval pipeline exposing a /search endpoint. The _search_local method sends an HTTP POST to config.local_base_url (defaulting to http://localhost:8000), expecting a response that combines dense embeddings, sparse BM25 scores, and cross-encoder reranking.

Results are normalized into the standard SearchResult schema, preserving chunk IDs, confidence scores, and metadata. This allows teams to run containerized vector databases (like FAISS or Milvus) behind a custom API while maintaining the same agent integration code.

Dify Cloud (Managed Vector Store)

Best for: Rapid prototyping and managed SaaS workflows.

For teams using Dify's cloud platform, the _search_dify method injects the user-provided dify_api_key and optional dataset_id into an authenticated POST request. The implementation handles token management and rate limiting, decoding Dify's response into the common SearchResult format.

This pattern eliminates infrastructure management but requires internet connectivity and API key configuration via KnowledgeBaseConfig.

RAPTOR Tree Index (Hierarchical Retrieval)

Best for: Multi-level document summarization and long-context reasoning.

The RAPTOR (Recursive Abstractive Processing for Tree-Organized Retrieval) pattern implements hierarchical semantic indexing. The _search_raptor method posts queries to a RAPTOR service (configured via raptor_base_url), which traverses a tree-structured index to retrieve node-level summaries.

Each result maps to a SearchResult with a synthetic chunk_id encoding the node level (e.g., root/branch/leaf), enabling the agent to reference specific levels of abstraction in the source material.

GraphRAG Knowledge Graph (Entity-Based Retrieval)

Best for: Relational data, community detection, and entity-centric queries.

The GraphRAG pattern integrates graph-based retrieval through the _search_graphrag method. This implementation posts queries to a GraphRAG service and parses either entity or community results, depending on the query intent.

Unlike vector-based patterns that return text chunks, this pattern can return structured subgraphs or community summaries, which are normalized into the SearchResult payload with metadata indicating the result type (entity vs. community).

Architectural Flow

The integration follows a strict five-phase pipeline:

1. Configuration – chapter3/contextual-retrieval/config.py defines KnowledgeBaseConfig and the KnowledgeBaseType enum. Users select backends via CLI arguments (--kb-type) or YAML configuration.

2. Tool Instantiation – KnowledgeBaseTools performs lazy initialization (see _get_offline_retriever at lines 85–90), loading indices only when first queried.

3. Tool Call – The LLM emits a function call to knowledge_base_search. The dispatcher selects the concrete implementation (_search_offline, _search_local, etc.) based on the configured enum value.

4. Result Normalization – Each backend returns dictionaries conforming to the SearchResult dataclass (lines 14–31), ensuring the agent receives a uniform shape regardless of source.

5. Document Retrieval – For full-text access, the agent calls get_document, which routes to backend-specific implementations like _get_document_local or _get_document_dify.

Extending the Pattern

Adding a new retrieval service requires four steps:

1. Extend the enum – Add a value to KnowledgeBaseType in config.py.
2. Implement the handler – Create a _search_<new> method in KnowledgeBaseTools following the HTTP POST/GET pattern and JSON parsing logic.
3. Wire the dispatcher – Add an elif clause in knowledge_base_search to route to the new method.
4. Update configuration – Add CLI arguments in main.py and config file support for the new provider.

Because the LLM only interacts with the generic tool signatures, new backends become immediately available to the agent without modifying prompts or agent logic.

Code Examples

Example 1: Offline BM25 Engine

from chapter3.agentic_rag.tools import KnowledgeBaseTools, KnowledgeBaseConfig, KnowledgeBaseType

cfg = KnowledgeBaseConfig(
    type=KnowledgeBaseType.OFFLINE,
    offline_corpus_path="data/corpus.txt",
    offline_top_k=5,
    document_store_path="data/store.json"
)

kb = KnowledgeBaseTools(cfg)
results = kb.knowledge_base_search("How does a transformer attend to tokens?")
for r in results:
    print(f"[{r['score']:.2f}] {r['text'][:120]}…")

Example 2: CLI Switch to RAPTOR Service

python -m chapter3.agentic_rag.main \
    --kb-type raptor \
    --raptor-base-url http://localhost:9000 \
    --raptor-top-k 10

The KnowledgeBaseTools object initializes with the RAPTOR dispatcher; subsequent knowledge_base_search calls automatically route to _search_raptor.

Example 3: Agent-Side Tool Definition

from chapter3.agentic_rag.tools import get_tool_definitions
import openai

# Pass tool definitions to the LLM client

client = openai.ChatCompletion.create(
    model="gpt-4o",
    messages=[...],
    functions=get_tool_definitions(),
)

The LLM can then emit function calls like:

{
  "name": "knowledge_base_search",
  "arguments": {"query": "What are the safety guidelines for using a fire extinguisher?"}
}

The framework routes the call to the active backend and returns normalized results.

Summary

• The modular RAG architecture in bojieli/ai-agent-book decouples retrieval backends from agent logic through a unified tool interface.
• Five patterns are implemented: Offline BM25, Local Hybrid API, Dify Cloud, RAPTOR Tree Index, and GraphRAG Knowledge Graph.
• All backends conform to the SearchResult dataclass (lines 14–31 in tools.py), enabling interchangeable use.
• Lazy initialization in KnowledgeBaseTools ensures resources are loaded only when queried.
• New backends require only enum extension, method implementation, and dispatcher wiring—no agent code changes needed.

Frequently Asked Questions

How does the agent decide which RAG backend to use?

The backend selection happens at startup via KnowledgeBaseConfig in chapter3/contextual-retrieval/config.py. The configuration object is passed to KnowledgeBaseTools, which stores the KnowledgeBaseType enum value. When knowledge_base_search is called, an internal if/elif chain routes to the appropriate private method (e.g., _search_offline or _search_graphrag). The LLM itself is unaware of the specific backend; it only sees the generic tool signature.

Can I combine multiple RAG patterns in a single agent session?

The current architecture in tools.py dispatches to a single backend per KnowledgeBaseTools instance. To combine patterns, instantiate multiple KnowledgeBaseTools objects with different configurations, or extend the dispatcher to fan out queries to multiple backends and merge results before returning them to the SearchResult list.

What is the performance difference between Offline BM25 and Local Hybrid API?

Offline BM25 operates entirely in-memory with no network overhead, making it suitable for low-latency, small-scale scenarios. The Local Hybrid API incurs HTTP round-trip costs but supports dense retrieval and reranking, typically yielding higher relevance scores for semantic queries. The choice depends on whether your priority is offline capability (BM25) or retrieval quality (Hybrid).

How do I migrate from Dify Cloud to a self-hosted solution?

Migration requires changing the KnowledgeBaseConfig from type=KnowledgeBaseType.DIFY to KnowledgeBaseType.LOCAL (or OFFLINE), and updating the corresponding URL or corpus path. Because both backends return the same SearchResult structure, the agent prompts and business logic require no changes—only the configuration and data ingestion pipelines need updating.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}