102 lines
3.8 KiB
Python
102 lines
3.8 KiB
Python
# -*- 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
|
|
from agentscope.model import DashScopeChatModel
|
|
from agentscope.rag import QdrantStore, SimpleKnowledge, TextReader
|
|
from agentscope.tool import Toolkit
|
|
|
|
# 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",
|
|
),
|
|
)
|
|
|
|
|
|
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(
|
|
knowledge.retrieve_knowledge,
|
|
func_description=( # Provide a clear description for the tool
|
|
"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."
|
|
),
|
|
)
|
|
|
|
# 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())
|