examples/rag/integrate_with_agent.py

# -*- coding: utf-8 -*-
"""Wire :class:`RAGMiddleware` into an :class:`Agent` — library mode.

The middleware in :mod:`agentscope.middleware._rag` is the agent-side
half of RAG: given one or more :class:`~agentscope.rag.KnowledgeBase`
handles (each pairing an embedding model with a vector-store
collection), the middleware drives search on each new user turn and
feeds the matched chunks into the model context.

This example reuses the indexing pipeline shown in ``index_and_search.py``
(parse → chunk → embed → insert) and then attaches the same knowledge base
to two agents — one per mode:

- ``"static"``: on the first reasoning step of a reply, embed the
  user's question, search, and inject the top hits as a one-shot
  :class:`HintBlock` into the agent's context. The model sees the
  matched snippets but never "decides" to search.
- ``"agentic"`` (the default): expose a ``search_knowledge`` tool. The
  model decides when (and what) to search, the same way it decides any
  other tool call.

Run with::

    DASHSCOPE_API_KEY=sk-... python examples/rag/integrate_with_agent.py
"""
import asyncio
import os

from agentscope.agent import Agent
from agentscope.credential import DashScopeCredential
from agentscope.embedding import DashScopeEmbeddingModel
from agentscope.message import UserMsg
from agentscope.middleware import RAGMiddleware
from agentscope.model import DashScopeChatModel
from agentscope.rag import (
    ApproxTokenChunker,
    KnowledgeBase,
    QdrantStore,
    TextParser,
)
from agentscope.tool import Toolkit


COLLECTION = "demo-kb"

KNOWLEDGE: dict[str, bytes] = {
    "company-policy.md": (
        b"# Acme Remote Work Policy\n\n"
        b"Employees may work remotely up to three days per week. "
        b"Wednesdays are mandatory in-office days for the whole "
        b"engineering org so cross-team syncs land on a predictable "
        b"day.\n\n"
        b"Equipment stipend: each new hire receives a USD 1,500 "
        b"one-off stipend for a home-office setup. Receipts must be "
        b"submitted within 90 days of the start date.\n"
    ),
    "release-notes.md": (
        b"# AgentScope 3.0 release notes\n\n"
        b"- New ``agentscope.rag`` module: pluggable parser, chunker, "
        b"embedding, and vector-store backends.\n"
        b"- ``RAGMiddleware`` ships in two modes -- ``static`` for "
        b"automatic injection, ``agentic`` for tool-driven search.\n"
        b"- Knowledge base service supports embedded and dedicated "
        b"worker deployments through a single message-bus channel.\n"
    ),
}


async def index_corpus(knowledge: KnowledgeBase) -> None:
    """Index the demo corpus into the knowledge base.

    Identical pipeline to ``examples/rag/index_and_search.py`` — extracted
    as a helper here so the agent-side wiring stays the focus.  Each source
    file becomes one logical document; ``KnowledgeBase.insert_document``
    embeds and inserts every chunk in a single batch.
    """
    parser = TextParser()
    chunker = ApproxTokenChunker(chunk_size=256, overlap=32)
    for filename, file_bytes in KNOWLEDGE.items():
        sections = await parser.parse(file=file_bytes, filename=filename)
        chunks = await chunker.chunk(sections)
        await knowledge.insert_document(
            chunks,
            document_metadata={"filename": filename},
        )


def build_agent(
    name: str,
    *,
    chat_model: DashScopeChatModel,
    rag_mw: RAGMiddleware,
) -> Agent:
    """Construct an :class:`Agent` with the RAG middleware attached.

    The middleware is just one entry in the ``middlewares=`` list; it
    composes with every other middleware (tool offload, mem0, ...) the
    agent uses.
    """
    return Agent(
        name=name,
        system_prompt=(
            "You are a concise assistant. Use matched context when "
            "available; if you don't know, say so."
        ),
        model=chat_model,
        toolkit=Toolkit(),
        middlewares=[rag_mw],
    )


async def ask(agent: Agent, question: str) -> None:
    """Run one reply and print the final assistant message."""
    print(f"\n[{agent.name}] user: {question}")
    reply = await agent.reply(UserMsg(name="user", content=question))
    print(f"[{agent.name}] assistant: {reply.get_text_content()}")


async def main() -> None:
    """The main entry point of the example."""
    api_key = os.environ.get("DASHSCOPE_API_KEY")
    if not api_key:
        raise RuntimeError(
            "Set DASHSCOPE_API_KEY before running this example.",
        )

    credential = DashScopeCredential(api_key=api_key)
    chat_model = DashScopeChatModel(
        credential=credential,
        model="qwen-plus",
        stream=False,
    )
    embedding_model = DashScopeEmbeddingModel(
        credential=credential,
        model="text-embedding-v4",
        dimensions=1024,
    )

    store = QdrantStore(location=":memory:")
    async with store:
        # One :class:`KnowledgeBase` handle binds embedding + vector store +
        # collection together.  ``insert_document`` / ``search`` /
        # ``list_documents`` all go through it, and the backing
        # collection is created lazily on first use.
        knowledge = KnowledgeBase(
            name="acme-handbook",
            description="Acme HR policies and AgentScope 3.0 release notes.",
            embedding_model=embedding_model,
            vector_store=store,
            collection=COLLECTION,
        )

        await index_corpus(knowledge)

        # ---- Mode 1: static ----
        # Search is automatic on the first reasoning step. The injected
        # ``HintBlock`` is one-shot (removed after the model call) so it
        # doesn't poison the next turn.
        static_mw = RAGMiddleware(
            knowledge_bases=[knowledge],
            parameters=RAGMiddleware.Parameters(
                mode="static",
                top_k=3,
                emit_hint_event=False,
            ),
        )
        static_agent = build_agent(
            "rag-static-agent",
            chat_model=chat_model,
            rag_mw=static_mw,
        )
        await ask(
            static_agent,
            "How many remote days per week does Acme allow?",
        )

        # ---- Mode 2: agentic ----
        # The middleware exposes a ``search_knowledge`` tool instead of
        # auto-injecting. The model decides when to call it; it may
        # also pass ``knowledge_bases=[...]`` to scope the search when
        # multiple knowledge bases are bound.
        agentic_mw = RAGMiddleware(
            knowledge_bases=[knowledge],
            parameters=RAGMiddleware.Parameters(mode="agentic", top_k=3),
        )
        agentic_agent = build_agent(
            "rag-agentic-agent",
            chat_model=chat_model,
            rag_mw=agentic_mw,
        )
        await ask(
            agentic_agent,
            "Summarise what's new in the AgentScope 3.0 release notes.",
        )


if __name__ == "__main__":
    asyncio.run(main())