Integrating AgentScope with Alibaba Cloud Bailian Knowledge Studio

uv pip install "agentscope[full]" aiohttp

import os
import time
import hashlib
import requests

BASE_URL = "https://{workspaceId}.cn-beijing.maas.aliyuncs.com"
API_KEY = os.environ["DASHSCOPE_API_KEY"]
HEADERS = {"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}


def upload_file(file_path: str, category: str = "default") -> str:
    """Upload file to Data Center; return fileId."""
    file_size = str(os.path.getsize(file_path))
    file_name = os.path.basename(file_path)
    with open(file_path, "rb") as f:
        content_md5 = hashlib.md5(f.read()).hexdigest()

    # 1. Request upload lease
    resp = requests.post(
        f"{BASE_URL}/api/v1/connector/dash/applyFileUploadLease",
        headers=HEADERS,
        json={"category": category, "fileName": file_name,
              "sizeBytes": file_size, "contentMd5": content_md5})
    lease = resp.json()["data"]

    # 2. Upload to OSS
    with open(file_path, "rb") as f:
        requests.put(lease["param"]["url"],
                      headers=lease["param"]["headers"], data=f)

    # 3. Register file
    resp = requests.post(
        f"{BASE_URL}/api/v1/connector/dash/addFile",
        headers=HEADERS,
        json={"leaseId": lease["leaseId"], "category": category,
              "categoryType": "UNSTRUCTURED", "parser": "AUTO_SELECT"})
    return resp.json()["data"]["fileId"]


def create_kb_and_wait(file_ids: list[str], kb_name: str = "my-kb") -> str:
    """Create knowledge base, ingest files, poll until completion; return kb_id."""
    resp = requests.post(
        f"{BASE_URL}/api/v1/indices/rag/index/create_v2",
        headers=HEADERS,
        json={"name": kb_name, "structureType": "unstructured",
              "sinkType": "DEFAULT", "sourceType": "DATA_CENTER_FILE",
              "embeddingModelName": "text-embedding-v4", "chunkSize": 600,
              "docIds": file_ids,
              "dataSources": [{"sourceType": "DATA_CENTER_FILE"}]})
    data = resp.json()["data"]
    kb_id, job_id = data["pipelineId"], data["ingestionId"]

    while True:
        resp = requests.get(
            f"{BASE_URL}/api/v1/indices/rag/index_job/status",
            headers=HEADERS,
            params={"index_id": kb_id, "job_id": job_id})
        status = resp.json()["data"]["ingestion_status"]
        if status == "COMPLETED":
            break
        elif status in ("FAILED", "CANCELLED"):
            raise RuntimeError(f"Ingestion failed: {status}")
        time.sleep(3)
    return kb_id


# One-liner to build knowledge base
file_id = upload_file("product-guide.md")
kb_id = create_kb_and_wait([file_id])
print(f"Knowledge base ready: {kb_id}")

import aiohttp


class KnowledgeStudioRetriever:
    """Lightweight client for Knowledge Studio retrieval."""

    def __init__(self, api_key: str, base_url: str, agent_id: str):
        self.api_key = api_key
        self.base_url = base_url
        self.agent_id = agent_id  # Knowledge search service ID
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }

    async def search(self, query: str, top_n: int = 5) -> list[dict]:
        """Perform knowledge base search; return list of text chunks.

        Each chunk includes `score` and `text`.
        `top_n` truncates result count locally—the actual retrieval policy is set in the agent config.
        """
        url = f"{self.base_url}/api/v1/indices/knowledge/search"
        payload = {"agent_id": self.agent_id, "query": query}

        async with aiohttp.ClientSession() as session:
            async with session.post(url, headers=self.headers,
                                    json=payload) as resp:
                data = await resp.json()
                nodes = data["data"]["nodes"]
                return nodes[:top_n]

    def format_context(self, nodes: list[dict]) -> str:
        """Format retrieval results into prompt-friendly context text."""
        parts = []
        for n in nodes:
            doc_name = n.get("metadata", {}).get("doc_name", "Unknown source")
            parts.append(f"[{doc_name}]\n{n['text']}")
        return "\n---\n".join(parts)

from typing import Any
from agentscope.middleware import MiddlewareBase
from agentscope.agent import Agent
from agentscope.tool import ToolBase
from agentscope.tool import ToolChunk
from agentscope.message import TextBlock, ToolResultState
from agentscope.permission import (
    PermissionDecision, PermissionBehavior, PermissionContext,
)


class KnowledgeStudioRAGMiddleware(MiddlewareBase):
    """RAG middleware integrating Knowledge Studio into AgentScope agents.

    mode="static":  Retrieve before every inference and inject into system prompt  
    mode="agentic": Expose `search_knowledge` tool for autonomous model invocation
    """

    def __init__(self, retriever: KnowledgeStudioRetriever,
                 mode: str = "agentic", top_k: int = 5):
        self.retriever = retriever
        self.mode = mode
        self.top_k = top_k

    # ── static mode: on_system_prompt hook ───────────────────────

    async def on_system_prompt(self, agent: Agent,
                               current_prompt: str) -> str:
        """static mode: Inject retrieval results into system prompt.

        Called by AgentScope during ReAct reasoning step when assembling the system prompt
        (Transformer-type, serially chained).
        """
        if self.mode != "static":
            return current_prompt

        # Extract query from last user message in context
        recent_msgs = agent.state.context[-4:]
        user_query = ""
        for msg in reversed(recent_msgs):
            if msg.role == "user":
                user_query = msg.get_text_content()
                break
        if not user_query:
            return current_prompt

        # Retrieve and inject
        nodes = await self.retriever.search(user_query, self.top_k)
        if not nodes:
            return current_prompt

        context = self.retriever.format_context(nodes)
        return (f"{current_prompt}\n\n## Reference from Knowledge Base\n"
                f"Below are relevant excerpts—prioritize them in your response:\n\n{context}\n")

    # ── agentic mode: list_tools hook ───────────────────────────

    async def list_tools(self) -> list[ToolBase]:
        """agentic mode: Expose `search_knowledge` tool for model to invoke autonomously.

        AgentScope does NOT auto-call tools returned here.
        You must manually collect and pass them into the Toolkit when constructing the agent.
        """
        if self.mode != "agentic":
            return []

        retriever = self.retriever  # Capture via closure

        class SearchKnowledgeTool(ToolBase):
            """Search relevant text chunks in the knowledge base."""
            name: str = "search_knowledge"
            description: str = (
                "Search the knowledge base for relevant information. "
                "Use this when you need to look up facts, docs, or "
                "specific knowledge to answer the question.")
            input_schema: dict = {
                "type": "object",
                "properties": {
                    "query": {"type": "string",
                              "description": "The search query."},
                    "top_k": {"type": "integer", "default": 5,
                               "description": "Number of results to return."},
                },
                "required": ["query"],
            }
            is_concurrency_safe: bool = True
            is_read_only: bool = True

            async def check_permissions(
                self, tool_input: dict[str, Any],
                context: PermissionContext,
            ) -> PermissionDecision:
                return PermissionDecision(
                    behavior=PermissionBehavior.ALLOW,
                    message="auto-allow search")

            async def call(self, query: str, top_k: int = 5):
                """Perform knowledge base search and return relevant text chunks."""
                nodes = await retriever.search(query, top_k)
                if not nodes:
                    yield ToolChunk(
                        content=[TextBlock(text="No relevant results found.")],
                        state=ToolResultState.SUCCESS,
                        is_last=True,
                    )
                    return
                context = retriever.format_context(nodes)
                yield ToolChunk(
                    content=[TextBlock(text=context)],
                    state=ToolResultState.SUCCESS,
                    is_last=True,
                )

        return [SearchKnowledgeTool()]

import asyncio
import os
from agentscope.agent import Agent
from agentscope.model import DashScopeChatModel
from agentscope.credential import DashScopeCredential
from agentscope.tool import Toolkit
from agentscope.message import UserMsg

# 1. Instantiate DashScope Chat Model
chat_model = DashScopeChatModel(
    credential=DashScopeCredential(api_key=os.environ["DASHSCOPE_API_KEY"]),
    model="qwen-plus",
    stream=True,
)

# 2. Instantiate retrieval client
retriever = KnowledgeStudioRetriever(
    api_key=os.environ["DASHSCOPE_API_KEY"],
    base_url=BASE_URL,
    agent_id="aid-xxxxxxxx",  # Retrieved from Knowledge Search service in console
)

async def run_static():
    """static mode: Auto-retrieve and inject before every inference."""
    mw = KnowledgeStudioRAGMiddleware(retriever, mode="static", top_k=3)
    agent = Agent(
        name="kb-assistant",
        system_prompt=(
            "You are a knowledge base Q&A assistant. Answer questions using retrieved content."
            "If no relevant info exists in the knowledge base, say so."
        ),
        model=chat_model,
        middlewares=[mw],
    )

    result = await agent.reply(
        UserMsg(name="user", content="How do I call the Qwen API?")
    )
    print(result.get_text_content())

asyncio.run(run_static())

Steps to call the Qwen API (summarized from knowledge base):

1. Obtain API Key
   Go to Alibaba Cloud’s Large Model Platform Bailian, enable model calling services,
   and generate an API Key…

2. Install SDK
   pip install dashscope

3. Call API
   Use dashscope.Generation API to make requests…

async def run_agentic():
    """agentic mode: Model autonomously decides when and what to retrieve."""
    mw = KnowledgeStudioRAGMiddleware(retriever, mode="agentic", top_k=3)
    tools = await mw.list_tools()  # Manually collect tools exposed by middleware

    agent = Agent(
        name="kb-assistant",
        system_prompt=(
            "You are a knowledge base Q&A assistant. You may call the `search_knowledge` tool "
            "to retrieve from the knowledge base. Answer using retrieved results."
        ),
        model=chat_model,
        toolkit=Toolkit(tools=tools),
        middlewares=[mw],
    )

    result = await agent.reply(
        UserMsg(name="user", content="How do I call the Qwen API?")
    )
    print(result.get_text_content())

asyncio.run(run_agentic())

There are several ways to call the Qwen API:

### 1. OpenAI-Compatible Interface (Recommended — simple & universal)
For developers familiar with OpenAI SDK…

### 2. DashScope SDK
Use Alibaba Cloud’s native SDK…

### 3. Direct HTTP Calls
Use curl or any HTTP client…

static_mw = KnowledgeStudioRAGMiddleware(retriever, mode="static", top_k=3)
agentic_mw = KnowledgeStudioRAGMiddleware(retriever, mode="agentic", top_k=5)
tools = await agentic_mw.list_tools()

agent = Agent(
    name="kb-assistant",
    system_prompt="You are a knowledge base Q&A assistant.",
    model=chat_model,
    toolkit=Toolkit(tools=tools),
    middlewares=[static_mw, agentic_mw],
)

from agentscope.state import AgentState
from agentscope.app.storage import RedisStorage

USER_ID = "user_123"
AGENT_ID = "agent_456"
SESSION_ID = "session_789"

async def run_with_persistence():
    async with RedisStorage(host="localhost", port=6379) as storage:
        record = await storage.get_session(
            user_id=USER_ID, agent_id=AGENT_ID, session_id=SESSION_ID)
        state = record.state if record else AgentState()

        agent = Agent(
            name="kb-assistant",
            system_prompt="You are a knowledge base Q&A assistant.",
            model=chat_model,
            toolkit=Toolkit(tools=await mw.list_tools()),
            middlewares=[mw],
            state=state,
        )

        result = await agent.reply(
            UserMsg(name="user", content="Is the feature mentioned earlier still supported?")
        )

        if record:
            # Update existing session state
            await storage.update_session_state(
                user_id=USER_ID, agent_id=AGENT_ID,
                session_id=SESSION_ID, state=agent.state)
        else:
            # First session: create session record first
            from agentscope.app.storage import SessionConfig
            await storage.upsert_session(
                user_id=USER_ID, agent_id=AGENT_ID,
                config=SessionConfig(workspace_id="", name="kb-session"),
                state=agent.state, session_id=SESSION_ID)