chore: initialize sandbox and overwrite remote content
Some checks failed
Pre-commit / run (ubuntu-latest) (push) Has been cancelled
Deploy Sphinx documentation to Pages / build_en (ubuntu-latest, 3.10) (push) Has been cancelled
Deploy Sphinx documentation to Pages / build_zh (ubuntu-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.12) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.12) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.12) (push) Has been cancelled
Some checks failed
Pre-commit / run (ubuntu-latest) (push) Has been cancelled
Deploy Sphinx documentation to Pages / build_en (ubuntu-latest, 3.10) (push) Has been cancelled
Deploy Sphinx documentation to Pages / build_zh (ubuntu-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.12) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.12) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.12) (push) Has been cancelled
This commit is contained in:
codex-bot
33
examples/functionality/agent_skill/README.md
Normal file
33
examples/functionality/agent_skill/README.md
Normal file
@@ -0,0 +1,33 @@
|
||||
# Agent Skills in AgentScope
|
||||
|
||||
[Agent Skill](https://claude.com/blog/skills) is an approach proposed by
|
||||
Anthropic to improve agent capabilities on specific tasks.
|
||||
|
||||
In this example, we demonstrate how to integrate Agent Skills into an
|
||||
ReAct agent in AgentScope via the `toolkit.register_agent_skill` API.
|
||||
|
||||
Specifically, we prepare a demonstration skill that helps the agent to
|
||||
learn about the AgentScope framework itself in the `skill` directory.
|
||||
In `main.py`, we register this skill to the agent's toolkit, and ask it
|
||||
to answer questions about AgentScope.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Install the latest version of AgentScope to run this example:
|
||||
|
||||
```bash
|
||||
pip install agentscope --upgrade
|
||||
```
|
||||
|
||||
Then, run the example with:
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
|
||||
> Note:
|
||||
> - The example is built with DashScope chat model. If you want to change the model used in this example, don't
|
||||
> forget to change the formatter at the same time! The corresponding relationship between built-in models and
|
||||
> formatters are list in [our tutorial](https://doc.agentscope.io/tutorial/task_prompt.html#id1)
|
||||
> - For local models, ensure the model service (like Ollama) is running before starting the agent.
|
||||
|
||||
79
examples/functionality/agent_skill/main.py
Normal file
79
examples/functionality/agent_skill/main.py
Normal file
@@ -0,0 +1,79 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The main entry point of the agent skill example."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import (
|
||||
Toolkit,
|
||||
execute_shell_command,
|
||||
execute_python_code,
|
||||
view_text_file,
|
||||
)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry point for the ReAct agent example."""
|
||||
toolkit = Toolkit()
|
||||
|
||||
# To use agent skills, your agent must be equipped with text file viewing
|
||||
# tools.
|
||||
toolkit.register_tool_function(execute_shell_command)
|
||||
toolkit.register_tool_function(execute_python_code)
|
||||
toolkit.register_tool_function(view_text_file)
|
||||
|
||||
# Register the agent skill
|
||||
toolkit.register_agent_skill(
|
||||
"./skill/analyzing-agentscope-library",
|
||||
)
|
||||
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="""You are a helpful assistant named Friday.
|
||||
|
||||
# IMPORTANT
|
||||
- Don't make any assumptions. All your knowledge about AgentScope library must come from your equipped skills.
|
||||
""", # noqa: E501
|
||||
model=DashScopeChatModel(
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
model_name="qwen3-max",
|
||||
enable_thinking=False,
|
||||
stream=True,
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
)
|
||||
|
||||
# First, let's take a look at the agent's system prompt
|
||||
print("\033[1;32mAgent System Prompt:\033[0m")
|
||||
print(agent.sys_prompt)
|
||||
print("\n")
|
||||
|
||||
print(
|
||||
"\033[1;32mResponse to Question 'What skills do you have?':\033[0m",
|
||||
)
|
||||
# We prepare two questions
|
||||
await agent(
|
||||
Msg("user", "What skills do you have?", "user"),
|
||||
)
|
||||
|
||||
print(
|
||||
"\n\033[1;32mResponse to Question 'How to create my own tool function "
|
||||
"for the agent in agentscope?':\033[0m",
|
||||
)
|
||||
# The second question
|
||||
await agent(
|
||||
Msg(
|
||||
"user",
|
||||
"How to custom tool function for the agent in agentscope?",
|
||||
"user",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,78 @@
|
||||
---
|
||||
name: Analyzing AgentScope Library
|
||||
description: This skill provides a way to retrieve information from the AgentScope library for analysis and decision-making.
|
||||
---
|
||||
|
||||
# Analyzing AgentScope Library
|
||||
|
||||
## Overview
|
||||
|
||||
This guide covers the essential operations for retrieving and answering questions about the AgentScope library.
|
||||
If you need to answer questions regarding the AgentScope library, or look up specific information, functions/classes,
|
||||
examples or guidance, this skill will help you achieve that.
|
||||
|
||||
## Quick Start
|
||||
|
||||
The skill provides the following key scripts:
|
||||
|
||||
- Search for guidance in the AgentScope tutorial.
|
||||
- Search for official examples and recommended implementations provided by AgentScope.
|
||||
- A quick interface to view AgentScope's Python library by given a module name (e.g. agentscope), and return the module's submodules, classes, and functions.
|
||||
|
||||
When being asked an AgentScope-related question, you can follow the steps below to find the relevant information:
|
||||
|
||||
First decide which of the three scripts to use based on the user's question.
|
||||
- If user asks for "how to use" types of questions, use the "Search for Guidance" script to find the relevant tutorial
|
||||
- If user asks for "how to implement/build" types of questions, first search for relevant examples. If not found, then
|
||||
consider what functions are needed and search in the guide/tutorial
|
||||
- If user asks for "how to initialize" types of questions, first search for relevant tutorials. If not found, then
|
||||
consider to search for the corresponding modules, classes, or functions in the library.
|
||||
|
||||
|
||||
### Search for Examples
|
||||
|
||||
First ask for the user's permission to clone the agentscope GitHub repository if you haven't done so:
|
||||
|
||||
```bash
|
||||
git clone -b main https://github.com/agentscope-ai/agentscope
|
||||
```
|
||||
|
||||
In this repo, the `examples` folder contains various examples demonstrating how to use different features of the
|
||||
AgentScope library.
|
||||
They are organized in a tree structure by different functionalities. You should use shell command like `ls` or `cat` to
|
||||
navigate and view the examples. Avoid using `find` command to search for examples, as the name of the example
|
||||
files may not directly relate to the functionality being searched for.
|
||||
|
||||
### Search for Guidance
|
||||
|
||||
Similarly, first ensure you have cloned the agentscope GitHub repository.
|
||||
|
||||
The source agentscope tutorial is located in the `docs/tutorials` folder of the agentscope GitHub repository. It's
|
||||
organized by the different sections. To search for guidance, go to the `docs/tutorials` folder and view the tutorial
|
||||
files by shell command like `ls` or `cat`.
|
||||
|
||||
|
||||
### Search for Targeted Modules
|
||||
|
||||
First, ensure you have installed the agentscope library in your environment:
|
||||
|
||||
```bash
|
||||
pip list | grep agentscope
|
||||
```
|
||||
|
||||
If not installed, ask the user for permission to install it by command:
|
||||
|
||||
```bash
|
||||
pip install agentscope
|
||||
```
|
||||
|
||||
Then, run the following script to search for specific modules, classes, or functions. It's suggested to start with
|
||||
`agentscope` as the root module name, and then specify the submodule name you want to search for.
|
||||
|
||||
```bash
|
||||
python view_agentscope_module.py --module agentscope
|
||||
```
|
||||
|
||||
About detailed usage, please refer to the `./view_agentscope_module.py` script (located in the same folder as this
|
||||
SKILL.md file).
|
||||
|
||||
@@ -0,0 +1,306 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# pylint: skip-file
|
||||
"""Get the signatures of functions and classes in the agentscope library."""
|
||||
from typing import Literal, Callable
|
||||
|
||||
import agentscope
|
||||
import inspect
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
def get_class_signature(cls: type) -> str:
|
||||
"""Get the signature of a class.
|
||||
|
||||
Args:
|
||||
cls (`type`):
|
||||
A class object.
|
||||
|
||||
Returns:
|
||||
str: The signature of the class.
|
||||
"""
|
||||
# Obtain class name and docstring
|
||||
class_name = cls.__name__
|
||||
class_docstring = cls.__doc__ or ""
|
||||
|
||||
# Construct the class string
|
||||
class_str = f"class {class_name}:\n"
|
||||
if class_docstring:
|
||||
class_str += f' """{class_docstring}"""\n'
|
||||
|
||||
# Obtain the module of the class
|
||||
methods = []
|
||||
for name, method in inspect.getmembers(cls, predicate=inspect.isfunction):
|
||||
# Skip methods that are not part of the class
|
||||
if method.__qualname__.split(".")[0] != class_name:
|
||||
continue
|
||||
|
||||
if name.startswith("_") and name not in ["__init__", "__call__"]:
|
||||
continue
|
||||
|
||||
# Obtain the method's signature
|
||||
sig = inspect.signature(method)
|
||||
|
||||
# Construct the method string
|
||||
method_str = f" def {name}{sig}:\n"
|
||||
|
||||
# Add the method's docstring if it exists
|
||||
method_docstring = method.__doc__ or ""
|
||||
if method_docstring:
|
||||
method_str += f' """{method_docstring}"""\n'
|
||||
|
||||
methods.append(method_str)
|
||||
|
||||
class_str += "\n".join(methods)
|
||||
return class_str
|
||||
|
||||
|
||||
def get_function_signature(func: Callable) -> str:
|
||||
"""Get the signature of a function."""
|
||||
sig = inspect.signature(func)
|
||||
method_str = f"def {func.__name__}{sig}:\n"
|
||||
|
||||
method_docstring = func.__doc__ or ""
|
||||
if method_docstring:
|
||||
method_str += f' """{method_docstring}"""\n'
|
||||
|
||||
return method_str
|
||||
|
||||
|
||||
class FuncOrCls(BaseModel):
|
||||
"""The class records the module, signature, docstring, reference, and
|
||||
type"""
|
||||
|
||||
module: str
|
||||
"""The module of the function or class."""
|
||||
signature: str
|
||||
"""The signature of the function or class."""
|
||||
docstring: str
|
||||
"""The docstring of the function or class."""
|
||||
reference: str
|
||||
"""The reference to the source code of the function or class"""
|
||||
type: Literal["function", "class"]
|
||||
"""The type of the function or class, either 'function' or 'class'."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
module: str,
|
||||
signature: str,
|
||||
docstring: str,
|
||||
reference: str,
|
||||
# pylint: disable=redefined-builtin
|
||||
type: Literal["function", "class"],
|
||||
) -> None:
|
||||
"""Initialize the FuncOrCls instance."""
|
||||
super().__init__(
|
||||
module=module,
|
||||
signature=signature.strip(),
|
||||
docstring=docstring.strip(),
|
||||
reference=reference,
|
||||
type=type,
|
||||
)
|
||||
|
||||
|
||||
def _truncate_docstring(docstring: str, max_length: int = 200) -> str:
|
||||
"""Truncate the docstring to a maximum length.
|
||||
|
||||
Args:
|
||||
docstring (`str`):
|
||||
The docstring to truncate.
|
||||
max_length (`int`, *optional*, defaults to 200):
|
||||
The maximum length of the docstring.
|
||||
|
||||
Returns:
|
||||
`str`:
|
||||
The truncated docstring.
|
||||
"""
|
||||
if len(docstring) > max_length:
|
||||
return docstring[:max_length] + "..."
|
||||
return docstring
|
||||
|
||||
|
||||
def get_agentscope_module_signatures() -> list[FuncOrCls]:
|
||||
"""Get the signatures of functions and classes in the agentscope library.
|
||||
|
||||
Returns:
|
||||
`list[FuncOrCls]`:
|
||||
A list of FuncOrCls instances representing the functions and
|
||||
classes in the agentscope library.
|
||||
"""
|
||||
signatures = []
|
||||
for module in agentscope.__all__:
|
||||
as_module = getattr(agentscope, module)
|
||||
path_module = ".".join(["agentscope", module])
|
||||
|
||||
# Functions
|
||||
if inspect.isfunction(as_module):
|
||||
file = inspect.getfile(as_module)
|
||||
source_lines, start_line = inspect.getsourcelines(as_module)
|
||||
signatures.append(
|
||||
FuncOrCls(
|
||||
module=path_module,
|
||||
signature=get_function_signature(as_module),
|
||||
docstring=_truncate_docstring(as_module.__doc__ or ""),
|
||||
reference=f"{file}: {start_line}-"
|
||||
f"{start_line + len(source_lines)}",
|
||||
type="function",
|
||||
),
|
||||
)
|
||||
|
||||
else:
|
||||
if not hasattr(as_module, "__all__"):
|
||||
continue
|
||||
|
||||
# Modules with __all__ attribute
|
||||
for name in as_module.__all__:
|
||||
func_or_cls = getattr(as_module, name)
|
||||
path_func_or_cls = ".".join([path_module, name])
|
||||
|
||||
if inspect.isclass(func_or_cls):
|
||||
file = inspect.getfile(func_or_cls)
|
||||
source_lines, start_line = inspect.getsourcelines(
|
||||
func_or_cls,
|
||||
)
|
||||
signatures.append(
|
||||
FuncOrCls(
|
||||
module=path_func_or_cls,
|
||||
signature=get_class_signature(func_or_cls),
|
||||
docstring=_truncate_docstring(
|
||||
func_or_cls.__doc__ or "",
|
||||
),
|
||||
reference=(
|
||||
f"{file}: {start_line}-"
|
||||
f"{start_line + len(source_lines)}"
|
||||
),
|
||||
type="class",
|
||||
),
|
||||
)
|
||||
|
||||
elif inspect.isfunction(func_or_cls):
|
||||
file = inspect.getfile(func_or_cls)
|
||||
source_lines, start_line = inspect.getsourcelines(
|
||||
func_or_cls,
|
||||
)
|
||||
signatures.append(
|
||||
FuncOrCls(
|
||||
module=path_func_or_cls,
|
||||
signature=get_function_signature(func_or_cls),
|
||||
docstring=_truncate_docstring(
|
||||
func_or_cls.__doc__ or "",
|
||||
),
|
||||
reference=(
|
||||
f"{file}: {start_line}-"
|
||||
f"{start_line + len(source_lines)}"
|
||||
),
|
||||
type="function",
|
||||
),
|
||||
)
|
||||
|
||||
return signatures
|
||||
|
||||
|
||||
def view_agentscope_library(
|
||||
module: str,
|
||||
) -> str:
|
||||
"""View AgentScope's Python library by given a module name
|
||||
(e.g. agentscope), and return the module's submodules, classes, and
|
||||
functions. Given a class name, return the class's documentation, methods,
|
||||
and their signatures. Given a function name, return the function's
|
||||
documentation and signature. If you don't have any information about
|
||||
AgentScope library, try to use "agentscope" to view the available top
|
||||
modules.
|
||||
|
||||
Note this function only provide the module's brief information.
|
||||
For more information, you should view the source code.
|
||||
|
||||
Args:
|
||||
module (`str`):
|
||||
The module name to view, which should be a module path separated
|
||||
by dots (e.g. "agentscope.models"). It can refer to a module,
|
||||
a class, or a function.
|
||||
"""
|
||||
if not module.startswith("agentscope"):
|
||||
return (
|
||||
f"Module '{module}' is invalid. The input module should be "
|
||||
f"'agentscope' or submodule of 'agentscope.xxx.xxx' "
|
||||
f"(separated by dots)."
|
||||
)
|
||||
|
||||
agentscope_top_modules = {}
|
||||
for as_module in agentscope.__all__:
|
||||
if as_module in ["__version__", "logger"]:
|
||||
continue
|
||||
agentscope_top_modules[as_module] = getattr(
|
||||
agentscope,
|
||||
as_module,
|
||||
).__doc__
|
||||
|
||||
# top modules
|
||||
if module == "agentscope":
|
||||
top_modules_description = (
|
||||
[
|
||||
"The top-level modules in AgentScope library:",
|
||||
]
|
||||
+ [
|
||||
f"- agentscope.{k}: {v}"
|
||||
for k, v in agentscope_top_modules.items()
|
||||
]
|
||||
+ [
|
||||
"You can further view the classes/function within above "
|
||||
"modules by calling this function with the above module name.",
|
||||
]
|
||||
)
|
||||
return "\n".join(top_modules_description)
|
||||
|
||||
# class, functions
|
||||
modules = get_agentscope_module_signatures()
|
||||
for as_module in modules:
|
||||
if as_module.module == module:
|
||||
return f"""- The signature of '{module}':
|
||||
```python
|
||||
{as_module.signature}
|
||||
```
|
||||
|
||||
- Source code reference: {as_module.reference}"""
|
||||
|
||||
# two-level modules
|
||||
collected_modules = []
|
||||
for as_module in modules:
|
||||
if as_module.module.startswith(module):
|
||||
collected_modules.append(as_module)
|
||||
|
||||
if len(collected_modules) > 0:
|
||||
collected_modules_content = (
|
||||
[
|
||||
f"The classes/functions and their truncated docstring in "
|
||||
f"'{module}' module:",
|
||||
]
|
||||
+ [f"- {_.module}: {repr(_.docstring)}" for _ in collected_modules]
|
||||
+ [
|
||||
"The docstring is truncated for limited context. For detailed "
|
||||
"signature and methods, call this function with the above "
|
||||
"module name",
|
||||
]
|
||||
)
|
||||
|
||||
return "\n".join(collected_modules_content)
|
||||
|
||||
return (
|
||||
f"Module '{module}' not found. Use 'agentscope' to view the "
|
||||
f"top-level modules to ensure the given module is valid."
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--module",
|
||||
type=str,
|
||||
default="agentscope",
|
||||
help="The module name to view, e.g. 'agentscope'",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
res = view_agentscope_library(module=args.module)
|
||||
print(res)
|
||||
158
examples/functionality/long_term_memory/mem0/README.md
Normal file
158
examples/functionality/long_term_memory/mem0/README.md
Normal file
@@ -0,0 +1,158 @@
|
||||
# Mem0 Long-Term Memory in AgentScope
|
||||
|
||||
This example demonstrates how to
|
||||
|
||||
- use Mem0LongTermMemory to provide persistent semantic memory storage for AgentScope agents,
|
||||
- record and retrieve conversation history and user preferences across sessions,
|
||||
- integrate long-term memory with ReAct agents for context-aware conversations, and
|
||||
- configure DashScope embedding models and Qdrant vector store for memory management.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Python 3.10 or higher
|
||||
- DashScope API key from Alibaba Cloud
|
||||
|
||||
|
||||
## QuickStart
|
||||
|
||||
Install agentscope and ensure you have a valid DashScope API key in your environment variables.
|
||||
|
||||
> Note: The example is built with DashScope chat model and embedding model. If you want to use OpenAI models instead,
|
||||
> modify the model initialization in the example code accordingly.
|
||||
|
||||
```bash
|
||||
# Install agentscope from source
|
||||
cd {PATH_TO_AGENTSCOPE}
|
||||
pip install -e .
|
||||
# Install dependencies
|
||||
pip install mem0ai
|
||||
```
|
||||
|
||||
Set up your API key:
|
||||
|
||||
```bash
|
||||
export DASHSCOPE_API_KEY='YOUR_API_KEY'
|
||||
```
|
||||
|
||||
Run the example:
|
||||
|
||||
```bash
|
||||
python memory_example.py
|
||||
```
|
||||
|
||||
The example will:
|
||||
1. Initialize a Mem0LongTermMemory instance with DashScope models and Qdrant vector store
|
||||
2. Record a basic conversation to long-term memory
|
||||
3. Retrieve memories using semantic search
|
||||
4. Demonstrate ReAct agent integration with long-term memory for storing and retrieving user preferences
|
||||
|
||||
## Key Features
|
||||
|
||||
- **Vector-based Storage**: Uses Qdrant vector database for efficient semantic search and retrieval
|
||||
- **Flexible Configuration**: Support for multiple embedding models (OpenAI, DashScope) and vector stores
|
||||
- **Async Operations**: Full async support for non-blocking memory operations
|
||||
- **ReAct Agent Integration**: Seamless integration with AgentScope's ReActAgent and tool system
|
||||
|
||||
## Basic Usage
|
||||
|
||||
### Initialize Memory
|
||||
|
||||
```python
|
||||
import os
|
||||
from agentscope.memory import Mem0LongTermMemory
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.embedding import DashScopeTextEmbedding
|
||||
from mem0.vector_stores.configs import VectorStoreConfig
|
||||
|
||||
# Initialize with DashScope models and Qdrant vector store
|
||||
long_term_memory = Mem0LongTermMemory(
|
||||
agent_name="Friday",
|
||||
user_name="user_123",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen-max-latest",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY")
|
||||
),
|
||||
embedding_model=DashScopeTextEmbedding(
|
||||
model_name="text-embedding-v3",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
dimensions=1024
|
||||
),
|
||||
vector_store_config=VectorStoreConfig(
|
||||
provider="qdrant",
|
||||
config={
|
||||
"on_disk": True,
|
||||
"path": "./qdrant_data", # Your customized storage path
|
||||
"embedding_model_dims": 1024
|
||||
}
|
||||
)
|
||||
)
|
||||
```
|
||||
|
||||
> **Important**: If you change to a different embedding model or modify `embedding_model_dims`, you must either set a new storage path or delete the existing database files. Otherwise, a dimension mismatch error will occur.
|
||||
|
||||
### Integrate with ReAct Agent
|
||||
|
||||
```python
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.tool import Toolkit
|
||||
|
||||
# Create a ReAct agent with long-term memory
|
||||
toolkit = Toolkit()
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt=(
|
||||
"You are a helpful assistant named Friday. "
|
||||
"If you think there is relevant information about "
|
||||
"the user's preferences, you can record it to long-term "
|
||||
"memory using the tool `record_to_memory`. "
|
||||
"If you need to retrieve information from long-term "
|
||||
"memory, use the tool `retrieve_from_memory`."
|
||||
),
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen-max-latest",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY")
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
long_term_memory=long_term_memory,
|
||||
long_term_memory_mode="both"
|
||||
)
|
||||
|
||||
# Use the agent
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content="When I travel to Hangzhou, I prefer to stay in a homestay",
|
||||
name="user"
|
||||
)
|
||||
response = await agent(msg)
|
||||
```
|
||||
|
||||
## Advanced Configuration
|
||||
|
||||
You can customize the mem0 config by directly set :
|
||||
|
||||
```python
|
||||
long_term_memory = Mem0LongTermMemory(
|
||||
agent_name="Friday",
|
||||
user_name="user_123",
|
||||
mem0_config=your_mem0_config # Pass your custom mem0 configuration
|
||||
)
|
||||
```
|
||||
|
||||
For more configuration options, refer to the [mem0 documentation](https://github.com/mem0ai/mem0).
|
||||
|
||||
## What's in the Example
|
||||
|
||||
The `memory_example.py` file demonstrates:
|
||||
|
||||
1. **Basic Memory Recording**: Recording user conversations to long-term memory
|
||||
2. **Memory Retrieval**: Searching for stored memories using semantic similarity
|
||||
3. **ReAct Agent Integration**: Using long-term memory with ReAct agents to store and retrieve user preferences automatically
|
||||
|
||||
## Reference
|
||||
|
||||
- [mem0 Documentation](https://github.com/mem0ai/mem0)
|
||||
- [Qdrant Vector Database](https://qdrant.tech/)
|
||||
185
examples/functionality/long_term_memory/mem0/memory_example.py
Normal file
185
examples/functionality/long_term_memory/mem0/memory_example.py
Normal file
@@ -0,0 +1,185 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Memory example demonstrating long-term memory functionality with mem0.
|
||||
|
||||
This module provides examples of how to use the Mem0LongTermMemory class
|
||||
for recording and retrieving persistent memories.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from mem0.vector_stores.configs import VectorStoreConfig
|
||||
from agentscope.memory import Mem0LongTermMemory
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.embedding import DashScopeTextEmbedding
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import Toolkit
|
||||
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Run the memory examples."""
|
||||
# Initialize long term memory
|
||||
long_term_memory = Mem0LongTermMemory(
|
||||
agent_name="Friday",
|
||||
user_name="user_123",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen-max-latest",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
),
|
||||
embedding_model=DashScopeTextEmbedding(
|
||||
model_name="text-embedding-v3",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
dimensions=1024,
|
||||
),
|
||||
vector_store_config=VectorStoreConfig(
|
||||
provider="qdrant",
|
||||
config={
|
||||
"on_disk": True,
|
||||
"path": "../memory/qdrant_data", # Specify custom path
|
||||
"embedding_model_dims": 1024,
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
# If you want to also use graph memory in mem0,
|
||||
# the following is an example of using Neo4j graph store.
|
||||
# from mem0.configs.base import MemoryConfig
|
||||
# from mem0.graphs.configs import GraphStoreConfig
|
||||
# long_term_memory = Mem0LongTermMemory(
|
||||
# agent_name="Friday",
|
||||
# user_name="user_123",
|
||||
# embedding_model=DashScopeTextEmbedding(
|
||||
# model_name="text-embedding-v3",
|
||||
# api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
# dimensions=1024,
|
||||
# ),
|
||||
# model=DashScopeChatModel(
|
||||
# model_name="qwen-max-latest",
|
||||
# api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
# stream=False,
|
||||
# ),
|
||||
# vector_store_config=VectorStoreConfig(
|
||||
# provider="qdrant",
|
||||
# config={
|
||||
# "on_disk": True,
|
||||
# "path": "../memory/qdrant_data", # Specify custom path
|
||||
# "embedding_model_dims": 1024,
|
||||
# }),
|
||||
# mem0_config=MemoryConfig(
|
||||
# graph_store=GraphStoreConfig(
|
||||
# provider="neo4j",
|
||||
# config={
|
||||
# "url": os.environ.get("NEO4J_URL",
|
||||
# "neo4j://localhost:7687"),
|
||||
# "username": os.environ.get("NEO4J_USER", "neo4j"),
|
||||
# "password": os.environ.get("NEO4J_PASSWORD",
|
||||
# "12345678"),
|
||||
# "database": "neo4j",
|
||||
# },
|
||||
# ),
|
||||
# ),
|
||||
# )
|
||||
|
||||
print("=== Long Term Memory Examples with mem0 ===\n")
|
||||
|
||||
# Example 1: Basic conversation recording
|
||||
print("1. Basic Conversation Recording")
|
||||
print("-" * 40)
|
||||
results = await long_term_memory.record(
|
||||
msgs=[
|
||||
Msg(
|
||||
role="user",
|
||||
content="Please help me book a hotel, preferably homestay",
|
||||
name="user",
|
||||
),
|
||||
],
|
||||
)
|
||||
print(f"Recorded conversation: {results}\n")
|
||||
|
||||
# Example 2: Retrieving memories
|
||||
print("2. Retrieving Memories")
|
||||
print("-" * 40)
|
||||
print("Searching for weather-related memories...")
|
||||
weather_memories = await long_term_memory.retrieve(
|
||||
msg=[
|
||||
Msg(
|
||||
role="user",
|
||||
content="What's the weather like today?",
|
||||
name="user",
|
||||
),
|
||||
],
|
||||
)
|
||||
print(f"Retrieved weather memories: {weather_memories}\n")
|
||||
|
||||
print("Searching for user preference memories...")
|
||||
preference_memories = await long_term_memory.retrieve(
|
||||
msg=[
|
||||
Msg(
|
||||
role="user",
|
||||
content=(
|
||||
"I prefer temperatures in Celsius and wind speed in km/h"
|
||||
),
|
||||
name="user",
|
||||
),
|
||||
],
|
||||
)
|
||||
print(f"Retrieved preference memories: {preference_memories}\n")
|
||||
|
||||
# Example 3: ReActAgent with long term memory
|
||||
print("3. ReActAgent with long term memory")
|
||||
print("-" * 40)
|
||||
|
||||
toolkit = Toolkit()
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt=(
|
||||
"You are a helpful assistant named Friday. "
|
||||
"If you think there is relevant information about "
|
||||
"user's preference, you can record it to the long term "
|
||||
"memory by tool call `record_to_memory`. "
|
||||
"If you need to retrieve information from the long term "
|
||||
"memory, you can use the tool call `retrieve_from_memory`."
|
||||
),
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen-max-latest",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
long_term_memory=long_term_memory,
|
||||
long_term_memory_mode="both",
|
||||
)
|
||||
|
||||
await agent.memory.clear()
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content="When I travel to Hangzhou, I prefer to stay in a homestay",
|
||||
name="user",
|
||||
)
|
||||
msg = await agent(msg)
|
||||
print(f"ReActAgent response: {msg.get_text_content()}\n")
|
||||
|
||||
msg = Msg(role="user", content="what preference do I have?", name="user")
|
||||
msg = await agent(msg)
|
||||
print(f"ReActAgent response: {msg.get_text_content()}\n")
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content="I prefer to visit the West Lake",
|
||||
name="user",
|
||||
)
|
||||
msg = await agent(msg)
|
||||
print(f"ReActAgent response: {msg.get_text_content()}\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
609
examples/functionality/long_term_memory/reme/README.md
Normal file
609
examples/functionality/long_term_memory/reme/README.md
Normal file
@@ -0,0 +1,609 @@
|
||||
# ReMe Long-Term Memory in AgentScope
|
||||
|
||||
This example demonstrates how to:
|
||||
|
||||
- Use ReMe (Reflection Memory) to provide three specialized types of persistent memory storage for AgentScope agents
|
||||
- Record and retrieve personal information, task execution trajectories, and tool usage patterns across sessions
|
||||
- Integrate long-term memory with ReActAgent for context-aware conversations and continuous learning
|
||||
- Configure DashScope embedding models and vector stores for efficient memory management
|
||||
|
||||
## Overview
|
||||
|
||||
ReMe (Reflection Memory) provides three types of long-term memory for intelligent agents:
|
||||
|
||||
1. **Personal Memory** (`ReMePersonalLongTermMemory`) - Records and retrieves persistent personal information, preferences, and facts about users
|
||||
2. **Task Memory** (`ReMeTaskLongTermMemory`) - Learns from task execution trajectories and retrieves relevant past experiences for similar tasks
|
||||
3. **Tool Memory** (`ReMeToolLongTermMemory`) - Records tool execution results and generates usage guidelines to improve tool calling
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Python 3.12 or higher
|
||||
- DashScope API key from Alibaba Cloud (for the examples)
|
||||
|
||||
## QuickStart
|
||||
|
||||
### Installation
|
||||
|
||||
```bash
|
||||
# Install agentscope from source
|
||||
cd {PATH_TO_AGENTSCOPE}
|
||||
pip install -e .
|
||||
|
||||
# Install required dependencies
|
||||
pip install reme-ai python-dotenv
|
||||
```
|
||||
|
||||
### Setup
|
||||
|
||||
Set up your API key:
|
||||
|
||||
```bash
|
||||
export DASHSCOPE_API_KEY='YOUR_API_KEY'
|
||||
```
|
||||
|
||||
Or create a `.env` file:
|
||||
|
||||
```bash
|
||||
DASHSCOPE_API_KEY=YOUR_API_KEY
|
||||
```
|
||||
|
||||
### Run Examples
|
||||
|
||||
```bash
|
||||
# Personal Memory Example - 5 core interfaces
|
||||
python personal_memory_example.py
|
||||
|
||||
# Task Memory Example - 5 core interfaces
|
||||
python task_memory_example.py
|
||||
|
||||
# Tool Memory Example - Complete workflow with ReActAgent
|
||||
python tool_memory_example.py
|
||||
```
|
||||
|
||||
> **Note**: The examples use DashScope models by default. To use OpenAI or other models, modify the model initialization in the example code accordingly.
|
||||
|
||||
## Key Features
|
||||
|
||||
- **Three Specialized Memory Types**: Personal, Task, and Tool memory for different use cases
|
||||
- **Dual Interface Design**: Both tool functions (for agent calling) and direct methods (for programmatic use)
|
||||
- **Vector-based Retrieval**: Efficient semantic search using embedding models and vector stores
|
||||
- **Async-first Architecture**: Full async/await support for non-blocking operations
|
||||
- **ReActAgent Integration**: Seamless integration with AgentScope's ReActAgent and Toolkit
|
||||
- **Automatic Context Management**: Uses async context managers for proper resource handling
|
||||
|
||||
## Core Concepts
|
||||
|
||||
### Memory Types and Their Use Cases
|
||||
|
||||
| Memory Type | Purpose | When to Use |
|
||||
|------------|---------|-------------|
|
||||
| **Personal Memory** | Store user preferences, habits, and personal facts | User profiles, personalized assistants, long-term user context |
|
||||
| **Task Memory** | Learn from task execution trajectories | Problem-solving, debugging, repeated workflows, learning from past successes |
|
||||
| **Tool Memory** | Record tool usage patterns and generate guidelines | Tool-using agents, improving tool call accuracy, avoiding past errors |
|
||||
|
||||
### Interface Design
|
||||
|
||||
**Personal Memory** and **Task Memory** provide **5 core interfaces**:
|
||||
|
||||
1. **`record_to_memory()`** - Tool function for agents to record memories (returns `ToolResponse`)
|
||||
2. **`retrieve_from_memory()`** - Tool function for agents to retrieve memories (returns `ToolResponse`)
|
||||
3. **`record()`** - Direct method for programmatic recording (returns `None`)
|
||||
4. **`retrieve()`** - Direct method for programmatic retrieval (returns `str`)
|
||||
5. **ReActAgent Integration** - Use memory with `long_term_memory` and `long_term_memory_mode` parameters
|
||||
|
||||
**Tool Memory** provides **2 core interfaces** (no tool functions):
|
||||
|
||||
1. **`record()`** - Direct method for recording tool execution results (returns `None`)
|
||||
2. **`retrieve()`** - Direct method for retrieving tool usage guidelines (returns `str`)
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### 1. Personal Memory
|
||||
|
||||
**Use Case**: Record and retrieve user preferences, habits, and personal information.
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
import os
|
||||
from agentscope.memory import ReMePersonalLongTermMemory
|
||||
from agentscope.embedding import DashScopeTextEmbedding
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
|
||||
|
||||
async def main():
|
||||
# Initialize personal memory
|
||||
personal_memory = ReMePersonalLongTermMemory(
|
||||
agent_name="Friday",
|
||||
user_name="user_123",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen3-max",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
),
|
||||
embedding_model=DashScopeTextEmbedding(
|
||||
model_name="text-embedding-v4",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
dimensions=1024,
|
||||
),
|
||||
)
|
||||
|
||||
# Use async context manager (required!)
|
||||
async with personal_memory:
|
||||
# Interface 1: record_to_memory (tool function)
|
||||
result = await personal_memory.record_to_memory(
|
||||
thinking="User sharing travel preferences",
|
||||
content=[
|
||||
"I prefer to stay in homestays when traveling to Hangzhou",
|
||||
"I like to visit the West Lake in the morning",
|
||||
"I enjoy drinking Longjing tea",
|
||||
],
|
||||
)
|
||||
|
||||
# Interface 2: retrieve_from_memory (tool function)
|
||||
result = await personal_memory.retrieve_from_memory(
|
||||
keywords=["Hangzhou travel", "tea preference"],
|
||||
)
|
||||
|
||||
# Interface 3: record (direct method)
|
||||
await personal_memory.record(
|
||||
msgs=[
|
||||
Msg(role="user", content="I work as a software engineer", name="user"),
|
||||
Msg(role="assistant", content="Got it!", name="assistant"),
|
||||
],
|
||||
)
|
||||
|
||||
# Interface 4: retrieve (direct method)
|
||||
memories = await personal_memory.retrieve(
|
||||
msg=Msg(role="user", content="What do you know about my work?", name="user"),
|
||||
)
|
||||
print(memories)
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
**Integration with ReActAgent** (Interface 5):
|
||||
|
||||
```python
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.tool import Toolkit
|
||||
|
||||
async def use_with_agent():
|
||||
personal_memory = ReMePersonalLongTermMemory(...)
|
||||
|
||||
async with personal_memory:
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You are Friday with long-term memory. Always record user information and retrieve memories when needed.",
|
||||
model=DashScopeChatModel(...),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=Toolkit(),
|
||||
memory=InMemoryMemory(),
|
||||
long_term_memory=personal_memory, # Attach personal memory
|
||||
long_term_memory_mode="both", # Enable both record and retrieve tools
|
||||
)
|
||||
|
||||
# Agent can now use record_to_memory and retrieve_from_memory as tools
|
||||
msg = Msg(role="user", content="I prefer staying in homestays", name="user")
|
||||
response = await agent(msg)
|
||||
```
|
||||
|
||||
### 2. Task Memory
|
||||
|
||||
**Use Case**: Learn from task execution trajectories and retrieve relevant experiences.
|
||||
|
||||
```python
|
||||
from agentscope.memory import ReMeTaskLongTermMemory
|
||||
|
||||
|
||||
async def main():
|
||||
# Initialize task memory
|
||||
task_memory = ReMeTaskLongTermMemory(
|
||||
agent_name="TaskAssistant",
|
||||
user_name="task_workspace_123", # Acts as workspace_id
|
||||
model=DashScopeChatModel(...),
|
||||
embedding_model=DashScopeTextEmbedding(...),
|
||||
)
|
||||
|
||||
async with task_memory:
|
||||
# Interface 1: record_to_memory with score
|
||||
result = await task_memory.record_to_memory(
|
||||
thinking="Recording successful debugging approach",
|
||||
content=[
|
||||
"For API 404 errors: Check route definition, verify URL path, ensure correct port",
|
||||
"Always use linter to catch typos in route paths",
|
||||
],
|
||||
score=0.95, # High score for successful trajectory
|
||||
)
|
||||
|
||||
# Interface 2: retrieve_from_memory
|
||||
result = await task_memory.retrieve_from_memory(
|
||||
keywords=["debugging", "API errors"],
|
||||
)
|
||||
|
||||
# Interface 3: record with score in direct method
|
||||
await task_memory.record(
|
||||
msgs=[
|
||||
Msg(role="user", content="I'm getting a 404 error", name="user"),
|
||||
Msg(role="assistant", content="Let's check the route path...", name="assistant"),
|
||||
Msg(role="user", content="Found the typo!", name="user"),
|
||||
],
|
||||
score=0.95, # Optional score for this trajectory
|
||||
)
|
||||
|
||||
# Interface 4: retrieve (direct method)
|
||||
experiences = await task_memory.retrieve(
|
||||
msg=Msg(role="user", content="How to debug API errors?", name="user"),
|
||||
)
|
||||
print(experiences)
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
**Integration with ReActAgent** (Interface 5):
|
||||
|
||||
```python
|
||||
async def use_with_agent():
|
||||
task_memory = ReMeTaskLongTermMemory(...)
|
||||
|
||||
async with task_memory:
|
||||
agent = ReActAgent(
|
||||
name="TaskAssistant",
|
||||
sys_prompt="You are a task assistant. Record solutions and retrieve past experiences before solving problems.",
|
||||
model=DashScopeChatModel(...),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=Toolkit(),
|
||||
memory=InMemoryMemory(),
|
||||
long_term_memory=task_memory,
|
||||
long_term_memory_mode="both",
|
||||
)
|
||||
|
||||
# Agent learns from task executions over time
|
||||
msg = Msg(role="user", content="How should I optimize database queries?", name="user")
|
||||
response = await agent(msg)
|
||||
```
|
||||
|
||||
### 3. Tool Memory
|
||||
|
||||
**Use Case**: Record tool execution results and generate usage guidelines for better tool calling.
|
||||
|
||||
**Complete Workflow**:
|
||||
|
||||
```python
|
||||
import json
|
||||
from datetime import datetime
|
||||
from agentscope.memory import ReMeToolLongTermMemory
|
||||
from agentscope.tool import Toolkit, ToolResponse
|
||||
from agentscope.message import Msg, TextBlock
|
||||
|
||||
|
||||
# Step 1: Define tools
|
||||
async def web_search(query: str, max_results: int = 5) -> ToolResponse:
|
||||
"""Search the web for information."""
|
||||
result = f"Found {max_results} results for query: '{query}'"
|
||||
return ToolResponse(content=[TextBlock(type="text", text=result)])
|
||||
|
||||
|
||||
async def main():
|
||||
# Initialize tool memory
|
||||
tool_memory = ReMeToolLongTermMemory(
|
||||
agent_name="ToolBot",
|
||||
user_name="tool_workspace_demo",
|
||||
model=DashScopeChatModel(...),
|
||||
embedding_model=DashScopeTextEmbedding(...),
|
||||
)
|
||||
|
||||
async with tool_memory:
|
||||
# Step 2: Record tool execution history (accepts JSON strings in msgs)
|
||||
tool_result = {
|
||||
"create_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"tool_name": "web_search",
|
||||
"input": {"query": "Python asyncio tutorial", "max_results": 10},
|
||||
"output": "Found 10 results for query: 'Python asyncio tutorial'",
|
||||
"token_cost": 150,
|
||||
"success": True,
|
||||
"time_cost": 2.3
|
||||
}
|
||||
|
||||
# Interface 1: record (accepts JSON strings in message content)
|
||||
await tool_memory.record(
|
||||
msgs=[Msg(role="assistant", content=json.dumps(tool_result), name="assistant")],
|
||||
)
|
||||
|
||||
# Step 3: Retrieve tool guidelines
|
||||
# Interface 2: retrieve returns summarized guidelines
|
||||
guidelines = await tool_memory.retrieve(
|
||||
msg=Msg(role="user", content="web_search", name="user"),
|
||||
)
|
||||
|
||||
# Step 4: Inject guidelines into agent system prompt
|
||||
toolkit = Toolkit()
|
||||
toolkit.register_tool_function(web_search)
|
||||
|
||||
base_prompt = "You are ToolBot, a helpful AI assistant."
|
||||
enhanced_prompt = f"{base_prompt}\n\n# Tool Guidelines:\n{guidelines}"
|
||||
|
||||
agent = ReActAgent(
|
||||
name="ToolBot",
|
||||
sys_prompt=enhanced_prompt, # Guidelines enhance tool usage
|
||||
model=DashScopeChatModel(...),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
)
|
||||
|
||||
# Agent now uses tools with learned guidelines
|
||||
msg = Msg(role="user", content="Search for Python design patterns", name="user")
|
||||
response = await agent(msg)
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
> **Note**: Tool Memory does NOT provide `record_to_memory()` and `retrieve_from_memory()` tool functions. It only provides direct `record()` and `retrieve()` methods. Tool Memory is designed to be used programmatically to enhance agent system prompts, not as agent-callable tools.
|
||||
|
||||
## API Reference
|
||||
|
||||
### Common Parameters
|
||||
|
||||
All memory types share these initialization parameters:
|
||||
|
||||
```python
|
||||
ReMePersonalLongTermMemory(
|
||||
agent_name: str, # Name of the agent using this memory
|
||||
user_name: str, # User identifier (acts as workspace_id in ReMe)
|
||||
model: ModelWrapper, # LLM for summarization and processing
|
||||
embedding_model: EmbeddingWrapper, # Embedding model for vector retrieval
|
||||
vector_store_dir: str = "./memory_vector_store", # Storage location
|
||||
)
|
||||
```
|
||||
|
||||
### Interface Specifications
|
||||
|
||||
#### Personal Memory
|
||||
|
||||
| Interface | Type | Signature | Returns | Description |
|
||||
|-----------|------|-----------|---------|-------------|
|
||||
| `record_to_memory` | Tool Function | `(thinking: str, content: list[str])` | `ToolResponse` | Record personal information with reasoning |
|
||||
| `retrieve_from_memory` | Tool Function | `(keywords: list[str], limit: int = 3)` | `ToolResponse` | Retrieve memories by keywords |
|
||||
| `record` | Direct Method | `(msgs: list[Msg])` | `None` | Record message conversations |
|
||||
| `retrieve` | Direct Method | `(msg: Msg, top_k: int = 3)` | `str` | Query-based retrieval |
|
||||
|
||||
**Parameters**:
|
||||
- `thinking`: Reasoning about what to record
|
||||
- `content`: List of strings to remember
|
||||
- `keywords`: Search keywords
|
||||
- `limit`: Results per keyword (tool function, default: 3)
|
||||
- `top_k`: Total results to retrieve (direct method, default: 3)
|
||||
|
||||
#### Task Memory
|
||||
|
||||
| Interface | Type | Signature | Returns | Description |
|
||||
|-----------|------|-----------|---------|-------------|
|
||||
| `record_to_memory` | Tool Function | `(thinking: str, content: list[str], score: float = 1.0)` | `ToolResponse` | Record task trajectory with score |
|
||||
| `retrieve_from_memory` | Tool Function | `(keywords: list[str], top_k: int = 5)` | `ToolResponse` | Retrieve experiences by keywords |
|
||||
| `record` | Direct Method | `(msgs: list[Msg], score: float = 1.0)` | `None` | Record message conversations with score |
|
||||
| `retrieve` | Direct Method | `(msg: Msg, top_k: int = 5)` | `str` | Query-based experience retrieval |
|
||||
|
||||
**Parameters**:
|
||||
- `thinking`: Reasoning about the task execution
|
||||
- `content`: Task execution information and insights
|
||||
- `score`: Success score for the trajectory (0.0-1.0, default: 1.0)
|
||||
- `keywords`: Search keywords (e.g., task type, domain)
|
||||
- `top_k`: Number of results to retrieve (default: 5)
|
||||
|
||||
#### Tool Memory
|
||||
|
||||
| Interface | Type | Signature | Returns | Description |
|
||||
|-----------|------|-----------|---------|-------------|
|
||||
| `record` | Direct Method | `(msgs: list[Msg])` | `None` | Record tool results as messages (JSON format) |
|
||||
| `retrieve` | Direct Method | `(msg: Msg)` | `str` | Retrieve guidelines for tools |
|
||||
|
||||
**Parameters**:
|
||||
- `msgs`: List of messages where `content` contains JSON strings with tool execution metadata:
|
||||
- `create_time`: Timestamp (`"%Y-%m-%d %H:%M:%S"`)
|
||||
- `tool_name`: Tool identifier
|
||||
- `input`: Parameters used (dict)
|
||||
- `output`: Execution result (str)
|
||||
- `token_cost`: Token usage (int)
|
||||
- `success`: Execution status (bool)
|
||||
- `time_cost`: Duration in seconds (float)
|
||||
- `msg`: Message containing tool name to retrieve guidelines for
|
||||
- **Note**: Tool Memory does NOT provide tool functions (`record_to_memory` and `retrieve_from_memory`). It only provides direct methods for programmatic use.
|
||||
|
||||
### ReActAgent Integration Modes
|
||||
|
||||
When attaching **Personal Memory** or **Task Memory** to ReActAgent, use the `long_term_memory_mode` parameter:
|
||||
|
||||
```python
|
||||
agent = ReActAgent(
|
||||
name="Assistant",
|
||||
long_term_memory=memory, # ReMePersonalLongTermMemory or ReMeTaskLongTermMemory
|
||||
long_term_memory_mode="both", # Options: "record", "retrieve", "both"
|
||||
# ... other parameters
|
||||
)
|
||||
```
|
||||
|
||||
**Modes**:
|
||||
- `"record"`: Only adds `record_to_memory` tool to agent
|
||||
- `"retrieve"`: Only adds `retrieve_from_memory` tool to agent
|
||||
- `"both"`: Adds both tools (recommended for most use cases)
|
||||
|
||||
> **Note**: Tool Memory does NOT support ReActAgent integration with tool functions. Use Tool Memory programmatically to enhance system prompts as shown in the Tool Memory example.
|
||||
|
||||
### Async Context Manager (Required!)
|
||||
|
||||
All ReMe memory types **must** be used with async context managers:
|
||||
|
||||
```python
|
||||
async with long_term_memory:
|
||||
# All memory operations must be within this context
|
||||
await long_term_memory.record(msgs=[...])
|
||||
result = await long_term_memory.retrieve(msg=...)
|
||||
```
|
||||
|
||||
This ensures:
|
||||
- Proper initialization of the ReMe backend
|
||||
- Resource cleanup after operations
|
||||
- Vector store connection management
|
||||
|
||||
### Custom Configuration
|
||||
|
||||
```python
|
||||
from agentscope.memory import ReMePersonalLongTermMemory
|
||||
|
||||
# Custom storage location and models
|
||||
memory = ReMePersonalLongTermMemory(
|
||||
agent_name="Friday",
|
||||
user_name="user_123",
|
||||
model=your_custom_model, # Any AgentScope-compatible LLM
|
||||
embedding_model=your_embedding, # Any AgentScope-compatible embedding model
|
||||
vector_store_dir="./custom_path", # Custom storage directory
|
||||
)
|
||||
```
|
||||
|
||||
## Example Files Overview
|
||||
|
||||
### `personal_memory_example.py`
|
||||
|
||||
Demonstrates **5 core interfaces** for personal memory:
|
||||
|
||||
1. **`record_to_memory()`** - Record user preferences using tool function
|
||||
2. **`retrieve_from_memory()`** - Search memories by keywords using tool function
|
||||
3. **`record()`** - Direct recording of message conversations
|
||||
4. **`retrieve()`** - Direct query-based retrieval
|
||||
5. **ReActAgent Integration** - Agent autonomously uses memory tools
|
||||
|
||||
**Key Features**:
|
||||
- Recording travel preferences, work habits, and personal information
|
||||
- Keyword-based and query-based retrieval
|
||||
- System prompt guidelines for agent memory usage
|
||||
- Automatic memory tool calling by ReActAgent
|
||||
|
||||
### `task_memory_example.py`
|
||||
|
||||
Demonstrates **5 core interfaces** for task memory:
|
||||
|
||||
1. **`record_to_memory()`** - Record task experiences with scores
|
||||
2. **`retrieve_from_memory()`** - Retrieve relevant experiences by keywords
|
||||
3. **`record()`** - Direct recording with trajectory scores
|
||||
4. **`retrieve()`** - Direct experience retrieval
|
||||
5. **ReActAgent Integration** - Agent learns from past task executions
|
||||
|
||||
**Key Features**:
|
||||
- Recording project planning, debugging, and development experiences
|
||||
- Score-based trajectory evaluation (0.0-1.0)
|
||||
- Learning from successful and failed attempts
|
||||
- Continuous improvement through experience retrieval
|
||||
|
||||
### `tool_memory_example.py`
|
||||
|
||||
Demonstrates the **complete workflow** for tool memory:
|
||||
|
||||
1. **Mock tools** - Define and register tools to Toolkit
|
||||
2. **Record tool history** - Store execution results with metadata using `record()`
|
||||
3. **Retrieve guidelines** - Get summarized usage guidelines using `retrieve()`
|
||||
4. **Enhance agent prompt** - Inject guidelines into system prompt
|
||||
5. **Use ReActAgent** - Agent uses tools with learned guidelines
|
||||
|
||||
**Key Features**:
|
||||
- JSON-formatted tool execution recording via direct `record()` method
|
||||
- Automatic guideline generation through summarization
|
||||
- Multi-tool guideline retrieval via direct `retrieve()` method
|
||||
- System prompt enhancement for better tool usage
|
||||
- **Note**: Tool Memory does NOT provide agent-callable tool functions
|
||||
|
||||
## Architecture
|
||||
|
||||
### Inheritance Hierarchy
|
||||
|
||||
```
|
||||
ReMeLongTermMemoryBase (abstract base)
|
||||
├── ReMePersonalLongTermMemory
|
||||
├── ReMeTaskLongTermMemory
|
||||
└── ReMeToolLongTermMemory
|
||||
```
|
||||
|
||||
**`ReMeLongTermMemoryBase`** provides:
|
||||
- Integration with ReMe library's `ReMeApp`
|
||||
- Async context manager implementation
|
||||
- Common interface definitions
|
||||
- Vector store and embedding management
|
||||
|
||||
### Memory Storage
|
||||
|
||||
- **Location**: `./memory_vector_store/` (configurable)
|
||||
- **Isolation**: Each `user_name` maintains separate storage
|
||||
- **Persistence**: Memories persist across sessions
|
||||
- **Format**: Vector embeddings with metadata
|
||||
|
||||
## Best Practices
|
||||
|
||||
### 1. System Prompt Design
|
||||
|
||||
For agents with long-term memory, clearly specify when to record and retrieve:
|
||||
|
||||
```python
|
||||
sys_prompt = """
|
||||
You are an assistant with long-term memory.
|
||||
|
||||
Recording Guidelines:
|
||||
- Record when users share personal information, preferences, or important facts
|
||||
- Record successful task execution approaches and solutions
|
||||
- Record tool execution results with detailed metadata
|
||||
|
||||
Retrieval Guidelines:
|
||||
- ALWAYS retrieve before answering questions about past information
|
||||
- Retrieve when dealing with similar tasks to past executions
|
||||
- Check tool guidelines before using tools
|
||||
"""
|
||||
```
|
||||
|
||||
### 2. Score Assignment (Task Memory)
|
||||
|
||||
Use meaningful scores to prioritize experiences:
|
||||
|
||||
```python
|
||||
# Successful trajectory
|
||||
await task_memory.record_to_memory(..., score=0.95)
|
||||
|
||||
# Partially successful
|
||||
await task_memory.record_to_memory(..., score=0.6)
|
||||
|
||||
# Failed trajectory (still useful to learn from)
|
||||
await task_memory.record_to_memory(..., score=0.2)
|
||||
```
|
||||
|
||||
### 3. Tool Memory Workflow
|
||||
|
||||
Follow this pattern for tool memory:
|
||||
|
||||
```
|
||||
1. Execute tool → 2. Record result → 3. Trigger summarization → 4. Retrieve guidelines → 5. Use in agent
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
**Issue**: `RuntimeError: Memory not initialized`
|
||||
- **Solution**: Always use `async with memory:` context manager
|
||||
|
||||
**Issue**: No memories retrieved
|
||||
- **Solution**: Ensure you've recorded memories first and check `user_name` matches
|
||||
|
||||
**Issue**: Tool memory not generating guidelines
|
||||
- **Solution**: Record multiple tool executions to trigger summarization
|
||||
|
||||
**Issue**: Agent not using memory tools
|
||||
- **Solution**: Check `long_term_memory_mode="both"` and verify system prompt encourages memory usage
|
||||
|
||||
## References
|
||||
|
||||
- [ReMe Library](https://github.com/modelscope/ReMe) - Core memory implementation
|
||||
- [AgentScope Documentation](https://github.com/modelscope/agentscope) - Framework documentation
|
||||
- [DashScope API](https://dashscope.aliyun.com/) - Model API for examples
|
||||
@@ -0,0 +1,295 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Personal memory example demonstrating ReMe personal memory.
|
||||
|
||||
This module provides examples of how to use the ReMePersonalMemory
|
||||
class.
|
||||
|
||||
The example demonstrates 5 core interfaces:
|
||||
1. record_to_memory - Tool function for explicit memory recording
|
||||
2. retrieve_from_memory - Tool function for keyword-based retrieval
|
||||
3. record - Direct method for recording message conversations
|
||||
4. retrieve - Direct method for query-based retrieval
|
||||
5. ReActAgent integration - Using personal memory with ReActAgent
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.embedding import DashScopeTextEmbedding
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.memory import ReMePersonalLongTermMemory
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import ToolResponse, Toolkit
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
async def test_record_to_memory(
|
||||
memory: ReMePersonalLongTermMemory,
|
||||
) -> None:
|
||||
"""Test the record_to_memory tool function interface."""
|
||||
print("Interface 1: record_to_memory (Tool Function)")
|
||||
print("-" * 70)
|
||||
print("Purpose: Explicit memory recording with structured content")
|
||||
print("Test case: Recording user's travel preferences...")
|
||||
|
||||
result: ToolResponse = await memory.record_to_memory(
|
||||
thinking=("The user is sharing their travel preferences and habits"),
|
||||
content=[
|
||||
"I prefer to stay in homestays when traveling to Hangzhou",
|
||||
"I like to visit the West Lake in the morning",
|
||||
"I enjoy drinking Longjing tea",
|
||||
],
|
||||
)
|
||||
result_text = " ".join(
|
||||
block.get("text", "")
|
||||
for block in result.content
|
||||
if block.get("type") == "text"
|
||||
)
|
||||
print(f"✓ Result: {result_text}")
|
||||
print(
|
||||
f"✓ Status: {'Success' if 'Success' in result_text else 'Failed'}",
|
||||
)
|
||||
print()
|
||||
|
||||
|
||||
async def test_retrieve_from_memory(
|
||||
memory: ReMePersonalLongTermMemory,
|
||||
) -> None:
|
||||
"""Test the retrieve_from_memory tool function interface."""
|
||||
print("Interface 2: retrieve_from_memory (Tool Function)")
|
||||
print("-" * 70)
|
||||
print("Purpose: Keyword-based memory retrieval")
|
||||
print()
|
||||
|
||||
result = await memory.retrieve_from_memory(
|
||||
keywords=["Hangzhou travel", "tea preference"],
|
||||
)
|
||||
retrieved_text = " ".join(
|
||||
block.get("text", "")
|
||||
for block in result.content
|
||||
if block.get("type") == "text"
|
||||
)
|
||||
print("✓ Retrieved memories:")
|
||||
print(f"{retrieved_text}")
|
||||
print()
|
||||
|
||||
|
||||
async def test_record_direct(memory: ReMePersonalLongTermMemory) -> None:
|
||||
"""Test the direct record method interface."""
|
||||
print("Interface 3: record (Direct Recording)")
|
||||
print("-" * 70)
|
||||
print("Purpose: Direct recording of message conversations")
|
||||
print()
|
||||
print("Test case: Recording work preferences and habits...")
|
||||
|
||||
try:
|
||||
await memory.record(
|
||||
msgs=[
|
||||
Msg(
|
||||
role="user",
|
||||
content=(
|
||||
"I work as a software engineer and prefer "
|
||||
"remote work"
|
||||
),
|
||||
name="user",
|
||||
),
|
||||
Msg(
|
||||
role="assistant",
|
||||
content=(
|
||||
"Understood! You're a software engineer who "
|
||||
"values remote work flexibility."
|
||||
),
|
||||
name="assistant",
|
||||
),
|
||||
Msg(
|
||||
role="user",
|
||||
content=(
|
||||
"I usually start my day at 9 AM with a "
|
||||
"cup of coffee"
|
||||
),
|
||||
name="user",
|
||||
),
|
||||
],
|
||||
)
|
||||
print("✓ Status: Successfully recorded conversation messages")
|
||||
print(
|
||||
"✓ Messages recorded: 3 messages (user-assistant dialogue)",
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"✗ Status: Failed - {str(e)}")
|
||||
print()
|
||||
|
||||
|
||||
async def test_retrieve_direct(memory: ReMePersonalLongTermMemory) -> None:
|
||||
"""Test the direct retrieve method interface."""
|
||||
print("Interface 4: retrieve (Direct Retrieval)")
|
||||
print("-" * 70)
|
||||
print("Purpose: Query-based memory retrieval using messages")
|
||||
print()
|
||||
print(
|
||||
"Test case: Querying 'What do you know about my "
|
||||
"work preferences?'...",
|
||||
)
|
||||
|
||||
memories = await memory.retrieve(
|
||||
msg=Msg(
|
||||
role="user",
|
||||
content="What do you know about my work preferences?",
|
||||
name="user",
|
||||
),
|
||||
)
|
||||
print("✓ Retrieved memories:")
|
||||
print(f"{memories if memories else 'No memories found'}")
|
||||
status = (
|
||||
"Success - Found memories"
|
||||
if memories
|
||||
else "No relevant memories found"
|
||||
)
|
||||
print(f"✓ Status: {status}")
|
||||
print()
|
||||
|
||||
|
||||
async def test_react_agent_with_memory(
|
||||
memory: ReMePersonalLongTermMemory,
|
||||
) -> None:
|
||||
"""Test ReActAgent integration with personal memory."""
|
||||
print("Interface 5: ReActAgent with Personal Memory")
|
||||
print("-" * 70)
|
||||
print(
|
||||
"Purpose: Demonstrate how ReActAgent uses personal memory tools",
|
||||
)
|
||||
print()
|
||||
print("Test case: Agent-driven memory recording and retrieval...")
|
||||
|
||||
toolkit = Toolkit()
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt=(
|
||||
"You are a helpful assistant named Friday with long-term "
|
||||
"memory capabilities. "
|
||||
"\n\n## Memory Management Guidelines:\n"
|
||||
"1. **Recording Memories**: When users share personal "
|
||||
"information, preferences, "
|
||||
"habits, or facts about themselves, ALWAYS record them "
|
||||
"using `record_to_memory` "
|
||||
"for future reference.\n"
|
||||
"\n2. **Retrieving Memories**: BEFORE answering questions "
|
||||
"about the user's preferences, "
|
||||
"past information, or personal details, you MUST FIRST "
|
||||
"call `retrieve_from_memory` "
|
||||
"to check if you have any relevant stored information. "
|
||||
"Do NOT rely solely on the "
|
||||
"current conversation context.\n"
|
||||
"\n3. **When to Retrieve**: Call `retrieve_from_memory` "
|
||||
"when:\n"
|
||||
" - User asks questions like 'what do I like?', "
|
||||
"'what are my preferences?', "
|
||||
"'what do you know about me?'\n"
|
||||
" - User asks about their past behaviors, habits, or "
|
||||
"preferences\n"
|
||||
" - User refers to information they mentioned before\n"
|
||||
" - You need context about the user to provide "
|
||||
"personalized responses\n"
|
||||
"\nAlways check your memory first before claiming you "
|
||||
"don't know something about the user."
|
||||
),
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen3-max",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
long_term_memory=memory,
|
||||
long_term_memory_mode="both",
|
||||
)
|
||||
|
||||
await agent.memory.clear()
|
||||
|
||||
print(
|
||||
"→ User: 'When I travel to Hangzhou, I prefer to stay in "
|
||||
"a homestay'",
|
||||
)
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content=(
|
||||
"When I travel to Hangzhou, I prefer to stay in " "a homestay"
|
||||
),
|
||||
name="user",
|
||||
)
|
||||
msg = await agent(msg)
|
||||
print(f"✓ Agent response: {msg.get_text_content()}\n")
|
||||
|
||||
print("→ User: 'what preference do I have?'")
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content="what preference do I have?",
|
||||
name="user",
|
||||
)
|
||||
msg = await agent(msg)
|
||||
print(f"✓ Agent response: {msg.get_text_content()}\n")
|
||||
|
||||
print(
|
||||
"✓ Status: Successfully demonstrated ReActAgent with "
|
||||
"personal memory",
|
||||
)
|
||||
print()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Demonstrate the 5 core interfaces of ReMePersonalMemory.
|
||||
|
||||
This example shows how to use:
|
||||
1. record_to_memory - Tool function for explicit memory recording
|
||||
2. retrieve_from_memory - Tool function for keyword-based retrieval
|
||||
3. record - Direct method for recording message conversations
|
||||
4. retrieve - Direct method for query-based retrieval
|
||||
5. ReActAgent integration - Using personal memory with ReActAgent
|
||||
"""
|
||||
long_term_memory = ReMePersonalLongTermMemory(
|
||||
agent_name="Friday",
|
||||
user_name="user_123",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen3-max",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
),
|
||||
embedding_model=DashScopeTextEmbedding(
|
||||
model_name="text-embedding-v4",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
dimensions=1024,
|
||||
),
|
||||
)
|
||||
|
||||
print("=" * 70)
|
||||
print("ReMePersonalMemory - Testing 5 Core Interfaces")
|
||||
print("=" * 70)
|
||||
print()
|
||||
|
||||
# Use async context manager to ensure proper initialization
|
||||
async with long_term_memory:
|
||||
# await test_record_to_memory(long_term_memory)
|
||||
# await test_retrieve_from_memory(long_term_memory)
|
||||
# await test_record_direct(long_term_memory)
|
||||
# await test_retrieve_direct(long_term_memory)
|
||||
await test_react_agent_with_memory(long_term_memory)
|
||||
|
||||
# Alternative way: manually call __aenter__ and __aexit__
|
||||
# This is equivalent to using "async with long_term_memory" above
|
||||
# await long_term_memory.__aenter__()
|
||||
# await test_react_agent_with_memory(long_term_memory)
|
||||
# await long_term_memory.__aexit__()
|
||||
|
||||
print("=" * 70)
|
||||
print("Testing Complete: All 5 Core Interfaces Verified!")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,342 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Task memory example demonstrating ReMe task memory.
|
||||
|
||||
This module provides examples of how to use the ReMeTaskMemory class
|
||||
using the ReMe library.
|
||||
|
||||
The example demonstrates 5 core interfaces:
|
||||
1. record_to_memory - Tool function for recording task information
|
||||
2. retrieve_from_memory - Tool function for keyword-based retrieval
|
||||
3. record - Direct method for recording message conversations
|
||||
with scores
|
||||
4. retrieve - Direct method for retrieving task experiences
|
||||
5. ReActAgent integration - Using task memory with ReActAgent
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.embedding import DashScopeTextEmbedding
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.memory import ReMeTaskLongTermMemory
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import ToolResponse, Toolkit
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
async def test_record_to_memory(memory: ReMeTaskLongTermMemory) -> None:
|
||||
"""Test the record_to_memory tool function interface."""
|
||||
print("Interface 1: record_to_memory (Tool Function)")
|
||||
print("-" * 70)
|
||||
print(
|
||||
"Purpose: Record task execution information with thinking and content",
|
||||
)
|
||||
print()
|
||||
print("Test case: Recording project planning task information...")
|
||||
|
||||
result: ToolResponse = await memory.record_to_memory(
|
||||
thinking=(
|
||||
"Recording project planning best practices and "
|
||||
"development approach"
|
||||
),
|
||||
content=[
|
||||
"For web application projects, break down into phases: "
|
||||
"Requirements gathering, Design, Development, Testing, "
|
||||
"Deployment",
|
||||
"Development phase recommendations: Frontend (React), "
|
||||
"Backend (FastAPI), Database (PostgreSQL), Agile "
|
||||
"methodology with 2-week sprints",
|
||||
"Dependency management: Use npm for frontend and pip for "
|
||||
"Python backend, maintain requirements.txt and "
|
||||
"package.json files",
|
||||
],
|
||||
score=0.9, # Optional: score for this trajectory (1.0)
|
||||
)
|
||||
result_text = " ".join(
|
||||
block.get("text", "")
|
||||
for block in result.content
|
||||
if block.get("type") == "text"
|
||||
)
|
||||
print(f"✓ Result: {result_text}")
|
||||
print(
|
||||
f"✓ Status: {'Success' if 'Success' in result_text else 'Failed'}",
|
||||
)
|
||||
print()
|
||||
|
||||
|
||||
async def test_retrieve_from_memory(
|
||||
memory: ReMeTaskLongTermMemory,
|
||||
) -> None:
|
||||
"""Test the retrieve_from_memory tool function interface."""
|
||||
print("Interface 2: retrieve_from_memory (Tool Function)")
|
||||
print("-" * 70)
|
||||
print("Purpose: Keyword-based retrieval of task experiences")
|
||||
print()
|
||||
print(
|
||||
"Test case: Searching with keywords 'project planning', "
|
||||
"'development phase'...",
|
||||
)
|
||||
|
||||
result = await memory.retrieve_from_memory(
|
||||
keywords=["project planning", "development phase"],
|
||||
)
|
||||
retrieved_text = " ".join(
|
||||
block.get("text", "")
|
||||
for block in result.content
|
||||
if block.get("type") == "text"
|
||||
)
|
||||
print("✓ Retrieved experiences:")
|
||||
print(f"{retrieved_text}")
|
||||
has_experiences = (
|
||||
retrieved_text and "No task experiences found" not in retrieved_text
|
||||
)
|
||||
status = (
|
||||
"Success - Found experiences"
|
||||
if has_experiences
|
||||
else "No relevant experiences found"
|
||||
)
|
||||
print(f"✓ Status: {status}")
|
||||
print()
|
||||
|
||||
|
||||
async def test_record_direct(memory: ReMeTaskLongTermMemory) -> None:
|
||||
"""Test the direct record method interface."""
|
||||
print("Interface 3: record (Direct Recording)")
|
||||
print("-" * 70)
|
||||
print("Purpose: Direct recording of message conversations with scores")
|
||||
print()
|
||||
print("Test case: Recording debugging task conversation...")
|
||||
|
||||
try:
|
||||
await memory.record(
|
||||
msgs=[
|
||||
Msg(
|
||||
role="user",
|
||||
content="I'm getting a 404 error on my API endpoint",
|
||||
name="user",
|
||||
),
|
||||
Msg(
|
||||
role="assistant",
|
||||
content=(
|
||||
"Let's troubleshoot: 1) Check if the route is "
|
||||
"properly defined, 2) Verify the URL path, "
|
||||
"3) Ensure the server is running on the correct "
|
||||
"port"
|
||||
),
|
||||
name="assistant",
|
||||
),
|
||||
Msg(
|
||||
role="user",
|
||||
content="Found it! The route path had a typo.",
|
||||
name="user",
|
||||
),
|
||||
Msg(
|
||||
role="assistant",
|
||||
content=(
|
||||
"Great! Always double-check route paths and use "
|
||||
"a linter to catch typos early."
|
||||
),
|
||||
name="assistant",
|
||||
),
|
||||
],
|
||||
score=0.95, # Optional: score (default: 1.0)
|
||||
)
|
||||
print("✓ Status: Successfully recorded debugging trajectory")
|
||||
print("✓ Messages recorded: 4 messages with score 0.95")
|
||||
except Exception as e:
|
||||
print(f"✗ Status: Failed - {str(e)}")
|
||||
print()
|
||||
|
||||
|
||||
async def test_retrieve_direct(memory: ReMeTaskLongTermMemory) -> None:
|
||||
"""Test the direct retrieve method interface."""
|
||||
print("Interface 4: retrieve (Direct Retrieval)")
|
||||
print("-" * 70)
|
||||
print("Purpose: Query-based retrieval using messages")
|
||||
print()
|
||||
print("Test case: Querying 'How to debug API errors?'...")
|
||||
|
||||
memories = await memory.retrieve(
|
||||
msg=Msg(
|
||||
role="user",
|
||||
content=(
|
||||
"How should I approach debugging API errors in my "
|
||||
"application?"
|
||||
),
|
||||
name="user",
|
||||
),
|
||||
)
|
||||
print("✓ Retrieved experiences:")
|
||||
print(f"{memories if memories else 'No experiences found'}")
|
||||
status = (
|
||||
"Success - Found experiences"
|
||||
if memories
|
||||
else "No relevant experiences found"
|
||||
)
|
||||
print(f"✓ Status: {status}")
|
||||
print()
|
||||
|
||||
|
||||
async def test_react_agent_with_memory(
|
||||
memory: ReMeTaskLongTermMemory,
|
||||
) -> None:
|
||||
"""Test ReActAgent integration with task memory."""
|
||||
print("Interface 5: ReActAgent with Task Memory")
|
||||
print("-" * 70)
|
||||
print(
|
||||
"Purpose: Demonstrate how ReActAgent uses task memory tools",
|
||||
)
|
||||
print()
|
||||
print(
|
||||
"Test case: Agent-driven task experience recording and "
|
||||
"retrieval...",
|
||||
)
|
||||
|
||||
toolkit = Toolkit()
|
||||
agent = ReActAgent(
|
||||
name="TaskAssistant",
|
||||
sys_prompt=(
|
||||
"You are a helpful task assistant named TaskAssistant "
|
||||
"with long-term task memory. "
|
||||
"\n\n## Task Memory Management Guidelines:\n"
|
||||
"1. **Recording Task Experiences**: When you provide "
|
||||
"technical solutions, solve problems, "
|
||||
"or complete tasks, ALWAYS record the key insights using "
|
||||
"`record_to_memory`. Include:\n"
|
||||
" - Specific techniques and approaches used\n"
|
||||
" - Best practices and implementation details\n"
|
||||
" - Lessons learned and important considerations\n"
|
||||
" - Step-by-step procedures that worked well\n"
|
||||
"\n2. **Retrieving Past Experiences**: BEFORE solving a "
|
||||
"problem or answering technical "
|
||||
"questions, you MUST FIRST call `retrieve_from_memory` "
|
||||
"to check if you have relevant "
|
||||
"past experiences. This helps you:\n"
|
||||
" - Avoid repeating past mistakes\n"
|
||||
" - Leverage proven solutions\n"
|
||||
" - Provide more accurate and tested approaches\n"
|
||||
"\n3. **When to Retrieve**: Always retrieve when:\n"
|
||||
" - Asked about technical topics or problem-solving "
|
||||
"approaches\n"
|
||||
" - Asked to provide recommendations or best practices\n"
|
||||
" - Dealing with tasks similar to ones you may have "
|
||||
"handled before\n"
|
||||
" - User explicitly asks 'what do you know about...?' "
|
||||
"or 'have you seen this before?'\n"
|
||||
"\nAlways check your task memory first to provide the "
|
||||
"most informed responses."
|
||||
),
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen3-max",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
long_term_memory=memory,
|
||||
long_term_memory_mode="both",
|
||||
)
|
||||
|
||||
await agent.memory.clear()
|
||||
|
||||
print(
|
||||
"→ User: 'Here are some database optimization techniques "
|
||||
"I learned'",
|
||||
)
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content=(
|
||||
"I just learned some valuable database optimization "
|
||||
"techniques for slow queries: "
|
||||
"1) Add indexes on foreign keys and WHERE clause columns "
|
||||
"to speed up joins and filtering. "
|
||||
"2) Use table partitioning to divide large tables by "
|
||||
"date or category for faster queries. "
|
||||
"3) Implement query result caching with Redis to avoid "
|
||||
"repeated database hits. "
|
||||
"4) Optimize JOIN order - put smallest tables first to "
|
||||
"reduce intermediate result sets. "
|
||||
"5) Use EXPLAIN ANALYZE to identify bottlenecks and "
|
||||
"missing indexes. "
|
||||
"Please record these optimization techniques for future "
|
||||
"reference."
|
||||
),
|
||||
name="user",
|
||||
)
|
||||
msg = await agent(msg)
|
||||
print(f"✓ Agent response: {msg.get_text_content()}\n")
|
||||
|
||||
print(
|
||||
"→ User: 'What do you know about database optimization?'",
|
||||
)
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content=(
|
||||
"What do you know about database optimization? "
|
||||
"Can you retrieve any past experiences?"
|
||||
),
|
||||
name="user",
|
||||
)
|
||||
msg = await agent(msg)
|
||||
print(f"✓ Agent response: {msg.get_text_content()}\n")
|
||||
|
||||
print()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Demonstrate the 5 core interfaces of ReMeTaskMemory.
|
||||
|
||||
This example shows how to use:
|
||||
1. record_to_memory - Tool function for recording task information
|
||||
2. retrieve_from_memory - Tool function for keyword-based retrieval
|
||||
3. record - Direct method for recording message conversations with scores
|
||||
4. retrieve - Direct method for retrieving task experiences
|
||||
5. ReActAgent integration - Using task memory with ReActAgent
|
||||
"""
|
||||
long_term_memory = ReMeTaskLongTermMemory(
|
||||
agent_name="TaskAssistant",
|
||||
user_name="task_workspace_123",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen3-max",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
),
|
||||
embedding_model=DashScopeTextEmbedding(
|
||||
model_name="text-embedding-v4",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
dimensions=1024,
|
||||
),
|
||||
)
|
||||
|
||||
print("=" * 70)
|
||||
print("ReMeTaskMemory - Testing 5 Core Interfaces")
|
||||
print("=" * 70)
|
||||
print()
|
||||
|
||||
# Use async context manager to ensure proper initialization
|
||||
async with long_term_memory:
|
||||
# await test_record_to_memory(long_term_memory)
|
||||
# await test_retrieve_from_memory(long_term_memory)
|
||||
# await test_record_direct(long_term_memory)
|
||||
# await test_retrieve_direct(long_term_memory)
|
||||
await test_react_agent_with_memory(long_term_memory)
|
||||
|
||||
# Alternative way: manually call __aenter__ and __aexit__
|
||||
# This is equivalent to using "async with long_term_memory" above
|
||||
# await long_term_memory.__aenter__()
|
||||
# await test_react_agent_with_memory(long_term_memory)
|
||||
# await long_term_memory.__aexit__()
|
||||
|
||||
print("=" * 70)
|
||||
print("Testing Complete: All 5 Core Interfaces Verified!")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,436 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Tool memory example demonstrating ReMe tool memory with ReActAgent.
|
||||
|
||||
This module demonstrates the complete workflow:
|
||||
1. Mock a tool function and register it to Toolkit
|
||||
2. Record tool execution results to tool memory using record()
|
||||
3. Retrieve tool usage guidelines using retrieve()
|
||||
4. Inject guidelines into ReActAgent's system prompt
|
||||
5. Use ReActAgent with tool memory
|
||||
|
||||
This workflow helps LLMs learn from past tool usage patterns and
|
||||
improve their tool calling decisions over time.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.embedding import DashScopeTextEmbedding
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.memory import ReMeToolLongTermMemory
|
||||
from agentscope.message import Msg
|
||||
from agentscope.message import TextBlock
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import Toolkit, ToolResponse
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Step 1: Mock tool functions
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def web_search(query: str, max_results: int = 5) -> ToolResponse:
|
||||
"""Search the web for information.
|
||||
|
||||
Args:
|
||||
query: The search query string
|
||||
max_results: Maximum number of results to return
|
||||
|
||||
Returns:
|
||||
ToolResponse containing search results
|
||||
"""
|
||||
# Simulate web search
|
||||
result = f"Found {max_results} results for query: '{query}'"
|
||||
return ToolResponse(
|
||||
content=[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=result,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
async def calculate(expression: str) -> ToolResponse:
|
||||
"""Calculate a mathematical expression.
|
||||
|
||||
Args:
|
||||
expression: Mathematical expression to evaluate
|
||||
|
||||
Returns:
|
||||
ToolResponse containing calculation result
|
||||
"""
|
||||
try:
|
||||
# Simple calculation (in real scenario, use safer evaluation)
|
||||
result = eval(expression) # noqa: S307
|
||||
return ToolResponse(
|
||||
content=[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=f"Result: {result}",
|
||||
),
|
||||
],
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResponse(
|
||||
content=[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=f"Error calculating '{expression}': {str(e)}",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Step 2: Record tool execution history to tool memory
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def record_tool_history(
|
||||
tool_memory: ReMeToolLongTermMemory,
|
||||
) -> None:
|
||||
"""Record historical tool execution results to tool memory.
|
||||
|
||||
This simulates past tool usage that the agent can learn from.
|
||||
"""
|
||||
print("=" * 70)
|
||||
print("Step 1: Recording Tool Execution History to Memory")
|
||||
print("=" * 70)
|
||||
print()
|
||||
|
||||
# Record successful web_search examples
|
||||
web_search_histories = [
|
||||
{
|
||||
"create_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"tool_name": "web_search",
|
||||
"input": {
|
||||
"query": "Python asyncio tutorial",
|
||||
"max_results": 10,
|
||||
},
|
||||
"output": (
|
||||
"Found 10 results for query: 'Python asyncio tutorial'"
|
||||
),
|
||||
"token_cost": 150,
|
||||
"success": True,
|
||||
"time_cost": 2.3,
|
||||
},
|
||||
{
|
||||
"create_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"tool_name": "web_search",
|
||||
"input": {
|
||||
"query": "machine learning basics",
|
||||
"max_results": 5,
|
||||
},
|
||||
"output": ("Found 5 results for query: 'machine learning basics'"),
|
||||
"token_cost": 120,
|
||||
"success": True,
|
||||
"time_cost": 1.8,
|
||||
},
|
||||
]
|
||||
|
||||
# Record failed web_search example (empty query)
|
||||
web_search_fail = {
|
||||
"create_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"tool_name": "web_search",
|
||||
"input": {
|
||||
"query": "",
|
||||
"max_results": 5,
|
||||
},
|
||||
"output": "Error: Query cannot be empty",
|
||||
"token_cost": 20,
|
||||
"success": False,
|
||||
"time_cost": 0.1,
|
||||
}
|
||||
|
||||
# Record calculate examples
|
||||
calculate_histories = [
|
||||
{
|
||||
"create_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"tool_name": "calculate",
|
||||
"input": {
|
||||
"expression": "2 + 2",
|
||||
},
|
||||
"output": "Result: 4",
|
||||
"token_cost": 30,
|
||||
"success": True,
|
||||
"time_cost": 0.05,
|
||||
},
|
||||
{
|
||||
"create_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"tool_name": "calculate",
|
||||
"input": {
|
||||
"expression": "10 * 5 + 3",
|
||||
},
|
||||
"output": "Result: 53",
|
||||
"token_cost": 30,
|
||||
"success": True,
|
||||
"time_cost": 0.05,
|
||||
},
|
||||
]
|
||||
|
||||
# Record all histories
|
||||
all_histories = (
|
||||
web_search_histories + [web_search_fail] + calculate_histories
|
||||
)
|
||||
|
||||
print(f"Recording {len(all_histories)} tool execution histories...")
|
||||
await tool_memory.record(
|
||||
msgs=[
|
||||
Msg(
|
||||
role="assistant",
|
||||
content=json.dumps(history),
|
||||
name="assistant",
|
||||
)
|
||||
for history in all_histories
|
||||
],
|
||||
)
|
||||
print(f"✓ Successfully recorded {len(all_histories)} tool executions")
|
||||
print()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Step 3: Retrieve tool guidelines and create enhanced system prompt
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def retrieve_tool_guidelines(
|
||||
tool_memory: ReMeToolLongTermMemory,
|
||||
tool_names: list[str],
|
||||
) -> str:
|
||||
"""Retrieve tool usage guidelines from memory.
|
||||
|
||||
Args:
|
||||
tool_memory: The ReMeToolMemory instance
|
||||
tool_names: List of tool names to retrieve guidelines for
|
||||
|
||||
Returns:
|
||||
Combined guidelines text to be added to system prompt
|
||||
"""
|
||||
print("=" * 70)
|
||||
print("Step 2: Retrieving Tool Usage Guidelines from Memory")
|
||||
print("=" * 70)
|
||||
print()
|
||||
|
||||
all_guidelines = []
|
||||
|
||||
for tool_name in tool_names:
|
||||
print(f"Retrieving guidelines for '{tool_name}'...")
|
||||
guidelines = await tool_memory.retrieve(
|
||||
msg=Msg(
|
||||
role="user",
|
||||
content=tool_name,
|
||||
name="user",
|
||||
),
|
||||
)
|
||||
|
||||
if guidelines:
|
||||
all_guidelines.append(
|
||||
f"## Guidelines for {tool_name}:\n{guidelines}",
|
||||
)
|
||||
print(f"✓ Retrieved guidelines for '{tool_name}'")
|
||||
print(f" Preview: {guidelines}")
|
||||
else:
|
||||
print(
|
||||
f"✓ No guidelines found for '{tool_name}' " "(first time use)",
|
||||
)
|
||||
print()
|
||||
|
||||
if all_guidelines:
|
||||
combined_guidelines = "\n\n".join(all_guidelines)
|
||||
guidelines_prompt = f"""
|
||||
# Tool Usage Guidelines (from past experience)
|
||||
|
||||
{combined_guidelines}
|
||||
|
||||
Please follow these guidelines when using the tools.
|
||||
"""
|
||||
return guidelines_prompt
|
||||
else:
|
||||
return ""
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Step 4: Use ReActAgent with tool memory
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def use_react_agent_with_tool_memory(
|
||||
toolkit: Toolkit,
|
||||
tool_guidelines: str,
|
||||
) -> None:
|
||||
"""Create and use ReActAgent with tool memory guidelines.
|
||||
|
||||
Args:
|
||||
toolkit: The Toolkit with registered tools
|
||||
tool_guidelines: Retrieved tool usage guidelines
|
||||
"""
|
||||
print("=" * 70)
|
||||
print("Step 3: Using ReActAgent with Tool Memory")
|
||||
print("=" * 70)
|
||||
print()
|
||||
|
||||
# Create enhanced system prompt with tool guidelines
|
||||
base_sys_prompt = (
|
||||
"You are a helpful AI assistant named ToolBot.\n"
|
||||
"You have access to various tools to help users complete "
|
||||
"their tasks.\n"
|
||||
"Please use the tools appropriately based on the user's "
|
||||
"requests."
|
||||
)
|
||||
|
||||
if tool_guidelines:
|
||||
sys_prompt = f"{base_sys_prompt}\n{tool_guidelines}"
|
||||
print(
|
||||
"✓ System prompt enhanced with tool memory guidelines",
|
||||
)
|
||||
else:
|
||||
sys_prompt = base_sys_prompt
|
||||
print(
|
||||
"✓ Using base system prompt (no guidelines available)",
|
||||
)
|
||||
|
||||
print()
|
||||
|
||||
# Create ReActAgent
|
||||
agent = ReActAgent(
|
||||
name="ToolBot",
|
||||
sys_prompt=sys_prompt,
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen-max",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
max_iters=5,
|
||||
)
|
||||
|
||||
print("✓ ReActAgent created successfully")
|
||||
print()
|
||||
|
||||
# Test queries
|
||||
test_queries = [
|
||||
"Search the web for 'Python design patterns'",
|
||||
"Calculate 15 * 7 + 23",
|
||||
]
|
||||
|
||||
print("-" * 70)
|
||||
print("Testing ReActAgent with tool memory...")
|
||||
print("-" * 70)
|
||||
print()
|
||||
|
||||
for i, query in enumerate(test_queries, 1):
|
||||
print(f"Query {i}: {query}")
|
||||
print("-" * 70)
|
||||
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content=query,
|
||||
name="user",
|
||||
)
|
||||
|
||||
response = await agent(msg)
|
||||
print(f"Response: {response.get_text_content()}")
|
||||
print()
|
||||
print()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Demonstrate the workflow of using ReMeToolMemory with ReActAgent.
|
||||
|
||||
This example shows:
|
||||
1. Create mock tools and register them to Toolkit
|
||||
2. Record historical tool execution results to tool memory
|
||||
3. Retrieve tool usage guidelines from memory
|
||||
4. Inject guidelines into ReActAgent's system prompt
|
||||
5. Use ReActAgent to complete tasks with tool memory
|
||||
"""
|
||||
print("=" * 70)
|
||||
print("ReMeToolMemory + ReActAgent Integration Example")
|
||||
print("=" * 70)
|
||||
print()
|
||||
print("This workflow demonstrates:")
|
||||
print("1. Mock tools → Register to Toolkit")
|
||||
print("2. Record tool execution history → Tool Memory")
|
||||
print("3. Retrieve guidelines → Enhance system prompt")
|
||||
print("4. Use ReActAgent with tool memory")
|
||||
print()
|
||||
|
||||
# Initialize tool memory
|
||||
tool_memory = ReMeToolLongTermMemory(
|
||||
agent_name="ToolBot",
|
||||
user_name="tool_workspace_demo",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen-max",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
),
|
||||
embedding_model=DashScopeTextEmbedding(
|
||||
model_name="text-embedding-v4",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
dimensions=1024,
|
||||
),
|
||||
)
|
||||
|
||||
# Create and register tools to toolkit
|
||||
toolkit = Toolkit()
|
||||
toolkit.register_tool_function(web_search)
|
||||
toolkit.register_tool_function(calculate)
|
||||
print()
|
||||
|
||||
# Use async context manager for tool memory
|
||||
async with tool_memory:
|
||||
# Step 1: Record historical tool executions to memory
|
||||
await record_tool_history(tool_memory)
|
||||
|
||||
# Step 2: Retrieve tool usage guidelines
|
||||
tool_names = ["web_search", "calculate"]
|
||||
tool_guidelines = await retrieve_tool_guidelines(
|
||||
tool_memory,
|
||||
tool_names,
|
||||
)
|
||||
|
||||
# Step 3: Use ReActAgent with enhanced system prompt
|
||||
await use_react_agent_with_tool_memory(
|
||||
toolkit,
|
||||
tool_guidelines,
|
||||
)
|
||||
|
||||
# Alternative way: manually call __aenter__ and __aexit__
|
||||
# This is equivalent to using "async with tool_memory" above
|
||||
# await tool_memory.__aenter__()
|
||||
# tool_guidelines = await retrieve_tool_guidelines(tool_memory, tool_names)
|
||||
# await tool_memory.__aexit__()
|
||||
|
||||
print("=" * 70)
|
||||
print("Workflow Complete!")
|
||||
print("=" * 70)
|
||||
print()
|
||||
print("Summary:")
|
||||
print("✓ Mock tools created and registered to Toolkit")
|
||||
print("✓ Historical tool executions recorded to tool memory")
|
||||
print("✓ Tool usage guidelines retrieved from memory")
|
||||
print("✓ ReActAgent system prompt enhanced with guidelines")
|
||||
print(
|
||||
"✓ ReActAgent successfully used tools with memory guidance",
|
||||
)
|
||||
print()
|
||||
print("Benefits:")
|
||||
print("- Agent learns from past tool usage patterns")
|
||||
print("- Reduced errors by following proven guidelines")
|
||||
print("- Better tool parameter selection")
|
||||
print("- Improved success rate over time")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
59
examples/functionality/mcp/README.md
Normal file
59
examples/functionality/mcp/README.md
Normal file
@@ -0,0 +1,59 @@
|
||||
# MCP in AgentScope
|
||||
|
||||
This example demonstrates how to
|
||||
|
||||
- create MCP client with different transports (SSE and Streamable HTTP) and type (Stateless and Stateful),
|
||||
- register MCP tool functions and use them in a ReAct agent, and
|
||||
- get MCP tool function as a local callable object from the MCP client.
|
||||
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Python 3.10 or higher
|
||||
- DashScope API key from Alibaba Cloud
|
||||
|
||||
## Installation
|
||||
|
||||
### Install AgentScope
|
||||
|
||||
```bash
|
||||
# Install from source
|
||||
cd {PATH_TO_AGENTSCOPE}
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
## QuickStart
|
||||
|
||||
Install agentscope and ensure you have a valid DashScope API key in your environment variables.
|
||||
|
||||
> Note: The example is built with DashScope chat model. If you want to change the model in this example, don't forget
|
||||
> to change the formatter at the same time! The corresponding relationship between built-in models and formatters are
|
||||
> list in [our tutorial](https://doc.agentscope.io/tutorial/task_prompt.html#id1)
|
||||
|
||||
```bash
|
||||
pip install agentscope
|
||||
```
|
||||
|
||||
Start the MCP servers by the following commands in two separate terminals:
|
||||
|
||||
```bash
|
||||
# In one terminal, run:
|
||||
python mcp_add.py
|
||||
|
||||
# In another terminal, run:
|
||||
python mcp_multiply.py
|
||||
```
|
||||
|
||||
Two MCP servers will be started on `http://127.0.0.1:8001` (SSE server) and `http://127.0.0.1:8002` (streamable
|
||||
HTTP server).
|
||||
|
||||
After starting the MCP servers, you can run the agent example:
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
|
||||
The agent will:
|
||||
1. Register the MCP tools from the servers
|
||||
2. Use a ReAct agent to solve a calculation problem (multiplying two numbers and then adding another number)
|
||||
3. Return structured output with the final result
|
||||
110
examples/functionality/mcp/main.py
Normal file
110
examples/functionality/mcp/main.py
Normal file
@@ -0,0 +1,110 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Demo showcasing ReAct agent with MCP tools using different transports.
|
||||
|
||||
This example demonstrates:
|
||||
- Registering MCP tools with different transports (sse and streamable_http)
|
||||
- Using a ReAct agent with registered MCP tools
|
||||
- Getting structured output from the agent
|
||||
|
||||
Before running this demo, please execute:
|
||||
python mcp_servers.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.mcp import HttpStatelessClient, HttpStatefulClient
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import Toolkit
|
||||
|
||||
|
||||
class NumberResult(BaseModel):
|
||||
"""A simple number result model for structured output."""
|
||||
|
||||
result: int = Field(description="The result of the calculation")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry of the MCP example."""
|
||||
|
||||
toolkit = Toolkit()
|
||||
|
||||
# Create a stateful MCP client to connect to the SSE MCP server
|
||||
# note you can also use the stateless client
|
||||
add_mcp_client = HttpStatefulClient(
|
||||
name="add_mcp",
|
||||
transport="sse",
|
||||
url="http://127.0.0.1:8001/sse",
|
||||
)
|
||||
|
||||
# Create a stateless MCP client to connect to the StreamableHTTP MCP server
|
||||
# note you can also use the stateful client
|
||||
multiply_mcp_client = HttpStatelessClient(
|
||||
name="multiply_mcp",
|
||||
transport="streamable_http",
|
||||
url="http://127.0.0.1:8002/mcp",
|
||||
)
|
||||
|
||||
# The stateful client must be connected before using
|
||||
await add_mcp_client.connect()
|
||||
|
||||
# Register the MCP clients to the toolkit
|
||||
await toolkit.register_mcp_client(add_mcp_client)
|
||||
await toolkit.register_mcp_client(multiply_mcp_client)
|
||||
|
||||
# Initialize the agent
|
||||
agent = ReActAgent(
|
||||
name="Jarvis",
|
||||
sys_prompt="You're a helpful assistant named Jarvis.",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen-max",
|
||||
api_key=os.environ["DASHSCOPE_API_KEY"],
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
)
|
||||
|
||||
# Run the agent with a calculation task
|
||||
res = await agent(
|
||||
Msg(
|
||||
"user",
|
||||
"Calculate 2345 multiplied by 3456, then add 4567 to the result,"
|
||||
" what is the final outcome?",
|
||||
"user",
|
||||
),
|
||||
structured_model=NumberResult,
|
||||
)
|
||||
|
||||
print(
|
||||
"Structured Output:\n"
|
||||
"```\n"
|
||||
f"{json.dumps(res.metadata, indent=4, ensure_ascii=False)}\n"
|
||||
"```",
|
||||
)
|
||||
|
||||
# AgentScope also allows developers to obtain the MCP tool as a local
|
||||
# callable object, and use it directly.
|
||||
add_tool_function = await add_mcp_client.get_callable_function(
|
||||
"add",
|
||||
# If wrap the MCP tool result into the ToolResponse object in
|
||||
# AgentScope
|
||||
wrap_tool_result=True,
|
||||
)
|
||||
|
||||
# Call it manually
|
||||
manual_res = await add_tool_function(a=5, b=10)
|
||||
print("When manually calling the MCP tool function:")
|
||||
print(manual_res)
|
||||
|
||||
# The stateful client should be disconnected manually!
|
||||
await add_mcp_client.close()
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
16
examples/functionality/mcp/mcp_add.py
Normal file
16
examples/functionality/mcp/mcp_add.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""An SSE MCP server with a simple add tool function."""
|
||||
|
||||
from mcp.server import FastMCP
|
||||
|
||||
|
||||
mcp = FastMCP("Add", port=8001)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def add(a: int, b: int) -> int:
|
||||
"""Add two numbers."""
|
||||
return a + b
|
||||
|
||||
|
||||
mcp.run(transport="sse")
|
||||
16
examples/functionality/mcp/mcp_multiply.py
Normal file
16
examples/functionality/mcp/mcp_multiply.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""An SSE MCP server with a simple multiply tool function."""
|
||||
|
||||
from mcp.server import FastMCP
|
||||
|
||||
|
||||
mcp = FastMCP("Multiply", port=8002)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def multiply(c: int, d: int) -> int:
|
||||
"""Multiply two numbers."""
|
||||
return c * d
|
||||
|
||||
|
||||
mcp.run(transport="streamable-http")
|
||||
30
examples/functionality/plan/README.md
Normal file
30
examples/functionality/plan/README.md
Normal file
@@ -0,0 +1,30 @@
|
||||
# Plan with ReAct Agent
|
||||
|
||||
This example demonstrates how to use the plan module in AgentScope to make an agent create and manage a plan formally.
|
||||
|
||||
Specifically, we provide two examples: manual specification plan and Agent-managed plan.
|
||||
|
||||
## Manual Specification Plan
|
||||
|
||||
In this example, we first manually specify a plan for the agent to follow, then we let the agent execute the plan step by step.
|
||||
|
||||
To execute this example, run:
|
||||
|
||||
```bash
|
||||
python main_manual_plan.py
|
||||
```
|
||||
|
||||
## Agent-managed Plan
|
||||
|
||||
In this example, we let the agent create and manage its own plan.
|
||||
Specifically, we use a query "Review the recent changes in AgentScope GitHub repository over the past month."
|
||||
|
||||
To run the example, execute:
|
||||
|
||||
```bash
|
||||
python main_agent_managed_plan.py
|
||||
```
|
||||
|
||||
> Note: The example is built with DashScope chat model. If you want to change the model in this example, don't forget
|
||||
> to change the **formatter** at the same time! The corresponding relationship between built-in models and formatters
|
||||
> are list in [our tutorial](https://doc.agentscope.io/tutorial/task_prompt.html#id1)
|
||||
65
examples/functionality/plan/main_agent_managed_plan.py
Normal file
65
examples/functionality/plan/main_agent_managed_plan.py
Normal file
@@ -0,0 +1,65 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The main entry point of the plan example."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.agent import ReActAgent, UserAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.plan import PlanNotebook
|
||||
from agentscope.tool import (
|
||||
Toolkit,
|
||||
execute_shell_command,
|
||||
execute_python_code,
|
||||
write_text_file,
|
||||
insert_text_file,
|
||||
view_text_file,
|
||||
)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry point for the plan example."""
|
||||
toolkit = Toolkit()
|
||||
toolkit.register_tool_function(execute_shell_command)
|
||||
toolkit.register_tool_function(execute_python_code)
|
||||
toolkit.register_tool_function(write_text_file)
|
||||
toolkit.register_tool_function(insert_text_file)
|
||||
toolkit.register_tool_function(view_text_file)
|
||||
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="""You're a helpful assistant named Friday.
|
||||
|
||||
# Target
|
||||
Your target is to finish the given task with careful planning.
|
||||
|
||||
# Note
|
||||
- You can equip yourself with plan related tools to help you plan and execute the given task.
|
||||
- The resouces from search engines are not always correct, you should collect information from multiple sources and give the final answer after careful consideration.
|
||||
""", # noqa
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen3-max-preview",
|
||||
api_key=os.environ["DASHSCOPE_API_KEY"],
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
enable_meta_tool=True,
|
||||
plan_notebook=PlanNotebook(),
|
||||
)
|
||||
user = UserAgent(name="user")
|
||||
|
||||
msg = Msg(
|
||||
"user",
|
||||
"Review the recent changes in AgentScope GitHub repository "
|
||||
"over the past month.",
|
||||
"user",
|
||||
)
|
||||
while True:
|
||||
msg = await agent(msg)
|
||||
msg = await user(msg)
|
||||
if msg.get_text_content() == "exit":
|
||||
break
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
101
examples/functionality/plan/main_manual_plan.py
Normal file
101
examples/functionality/plan/main_manual_plan.py
Normal file
@@ -0,0 +1,101 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Manual specification plan example."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.agent import ReActAgent, UserAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.plan import PlanNotebook, SubTask
|
||||
from agentscope.tool import (
|
||||
Toolkit,
|
||||
execute_shell_command,
|
||||
execute_python_code,
|
||||
write_text_file,
|
||||
insert_text_file,
|
||||
view_text_file,
|
||||
)
|
||||
|
||||
plan_notebook = PlanNotebook()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry point for the manual plan example."""
|
||||
|
||||
# Create the plan manually
|
||||
await plan_notebook.create_plan(
|
||||
name="Comprehensive Report on AgentScope",
|
||||
description="Study the code of AgentScope and write a comprehensive "
|
||||
"report about this framework.",
|
||||
expected_outcome="A markdown format report summarizing the features, "
|
||||
"architecture, advantages/disadvantages, and "
|
||||
"potential improvements of AgentScope.",
|
||||
subtasks=[
|
||||
SubTask(
|
||||
name="Clone the repository",
|
||||
description="Clone the AgentScope GitHub repository from "
|
||||
"agentscope-ai/agentscope, and ensure it's the "
|
||||
"latest version.",
|
||||
expected_outcome="A local copy of the AgentScope repository.",
|
||||
),
|
||||
SubTask(
|
||||
name="View the documentation",
|
||||
description="View the documentation of AgentScope in the "
|
||||
"repository.",
|
||||
expected_outcome="A comprehensive understanding of the "
|
||||
"features and usage of AgentScope.",
|
||||
),
|
||||
SubTask(
|
||||
name="Study the code",
|
||||
description="Study the code of AgentScope, focusing on the "
|
||||
"core modules and their interactions.",
|
||||
expected_outcome="A deep understanding of the architecture "
|
||||
"and implementation of AgentScope.",
|
||||
),
|
||||
SubTask(
|
||||
name="Summarize the findings",
|
||||
description="Summarize the findings from the documentation "
|
||||
"and code study, and write a comprehensive report "
|
||||
"in markdown format.",
|
||||
expected_outcome="A markdown format report",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
# Add basic tools
|
||||
toolkit = Toolkit()
|
||||
toolkit.register_tool_function(execute_shell_command)
|
||||
toolkit.register_tool_function(execute_python_code)
|
||||
toolkit.register_tool_function(write_text_file)
|
||||
toolkit.register_tool_function(insert_text_file)
|
||||
toolkit.register_tool_function(view_text_file)
|
||||
|
||||
# Create the agent
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You're a helpful assistant named Friday. Your target is "
|
||||
"to finish the given task with careful planning.",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen3-max-preview",
|
||||
api_key=os.environ["DASHSCOPE_API_KEY"],
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
plan_notebook=plan_notebook,
|
||||
)
|
||||
user = UserAgent(name="user")
|
||||
|
||||
msg = Msg(
|
||||
"user",
|
||||
"Now start to finish the task by the given plan",
|
||||
"user",
|
||||
)
|
||||
while True:
|
||||
msg = await agent(msg)
|
||||
msg = await user(msg)
|
||||
if msg.get_text_content() == "exit":
|
||||
break
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
40
examples/functionality/rag/README.md
Normal file
40
examples/functionality/rag/README.md
Normal file
@@ -0,0 +1,40 @@
|
||||
# RAG in AgentScope
|
||||
|
||||
This example includes three scripts to demonstrate how to use Retrieval-Augmented Generation (RAG) in AgentScope:
|
||||
|
||||
- the basic usage of RAG module in AgentScope in ``basic_usage.py``,
|
||||
- a simple agentic use case of RAG in ``agentic_usage.py``, and
|
||||
- integrate RAG into ``ReActAgent`` class by retrieving input message(s) at the beginning of each reply in ``react_agent_integration.py``.
|
||||
- build multimodal RAG in ``multimodal_rag.py``.
|
||||
|
||||
> The agentic usage and static integration has their own advantages and limitations.
|
||||
> - The agentic usage requires more powerful LLMs to manage the retrieval process, but it's more flexible and the agent can adjust the retrieval strategy dynamically
|
||||
> - The static integration is more straightforward and easier to implement, but it's less flexible and the input message maybe not specific enough, leading to less relevant retrieval results.
|
||||
|
||||
> Note: The example is built with DashScope chat model. If you want to change the model in this example, don't forget
|
||||
> to change the formatter at the same time! The corresponding relationship between built-in models and formatters are
|
||||
> list in [our tutorial](https://doc.agentscope.io/tutorial/task_prompt.html#id1)
|
||||
|
||||
## Quick Start
|
||||
|
||||
Install the latest agentscope library from PyPI or source, then run the following command to run the example:
|
||||
|
||||
- the basic usage:
|
||||
```bash
|
||||
python basic_usage.py
|
||||
```
|
||||
|
||||
- the agentic usage:
|
||||
```bash
|
||||
python agentic_usage.py
|
||||
```
|
||||
|
||||
- the static integration:
|
||||
```bash
|
||||
python react_agent_integration.py
|
||||
```
|
||||
|
||||
- the multimodal RAG:
|
||||
```bash
|
||||
python multimodal_rag.py
|
||||
```
|
||||
101
examples/functionality/rag/agentic_usage.py
Normal file
101
examples/functionality/rag/agentic_usage.py
Normal file
@@ -0,0 +1,101 @@
|
||||
# -*- 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 SimpleKnowledge, QdrantStore, 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())
|
||||
79
examples/functionality/rag/basic_usage.py
Normal file
79
examples/functionality/rag/basic_usage.py
Normal file
@@ -0,0 +1,79 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The main entry point of the RAG example."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.embedding import DashScopeTextEmbedding
|
||||
from agentscope.rag import (
|
||||
TextReader,
|
||||
PDFReader,
|
||||
QdrantStore,
|
||||
SimpleKnowledge,
|
||||
)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry point of the RAG example."""
|
||||
|
||||
# Create readers with chunking arguments
|
||||
reader = TextReader(chunk_size=1024)
|
||||
pdf_reader = PDFReader(chunk_size=1024, split_by="sentence")
|
||||
|
||||
# Read documents
|
||||
documents = await reader(
|
||||
text="I'm Tony Stank, my password is 123456. My best friend is James "
|
||||
"Rhodes.",
|
||||
)
|
||||
|
||||
# Read a sample PDF file
|
||||
pdf_path = os.path.join(
|
||||
os.path.abspath(os.path.dirname(__file__)),
|
||||
"example.pdf",
|
||||
)
|
||||
pdf_documents = await pdf_reader(pdf_path=pdf_path)
|
||||
|
||||
# Create a knowledge base with Qdrant as the embedding store and
|
||||
# DashScope as the embedding model
|
||||
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",
|
||||
),
|
||||
)
|
||||
|
||||
# Insert documents into the knowledge base
|
||||
await knowledge.add_documents(documents + pdf_documents)
|
||||
|
||||
# Retrieve relevant documents based on a given query
|
||||
docs = await knowledge.retrieve(
|
||||
query="What is Tony Stank's password?",
|
||||
limit=3,
|
||||
score_threshold=0.7,
|
||||
)
|
||||
print("Q1: What is Tony Stank's password?")
|
||||
for doc in docs:
|
||||
print(
|
||||
f"Document ID: {doc.id}, Score: {doc.score}, "
|
||||
f"Content: {doc.metadata.content['text']}",
|
||||
)
|
||||
|
||||
# Retrieve documents from the PDF file based on a query
|
||||
docs = await knowledge.retrieve(
|
||||
query="climate change",
|
||||
limit=3,
|
||||
score_threshold=0.2,
|
||||
)
|
||||
print("\n\nQ2: climate change")
|
||||
for doc in docs:
|
||||
print(
|
||||
f"Document ID: {doc.id}, Score: {doc.score}, "
|
||||
f"Content: {repr(doc.metadata.content['text'])}",
|
||||
)
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
BIN
examples/functionality/rag/example.pdf
Normal file
BIN
examples/functionality/rag/example.pdf
Normal file
Binary file not shown.
72
examples/functionality/rag/multimodal_rag.py
Normal file
72
examples/functionality/rag/multimodal_rag.py
Normal file
@@ -0,0 +1,72 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The example of how to use multimodal RAG in AgentScope"""
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.embedding import DashScopeMultiModalEmbedding
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.rag import ImageReader, SimpleKnowledge, QdrantStore
|
||||
|
||||
|
||||
path_image = "./example.png"
|
||||
plt.figure(figsize=(8, 3))
|
||||
plt.text(0.5, 0.5, "My name is Ming Li", ha="center", va="center", fontsize=30)
|
||||
plt.axis("off")
|
||||
plt.savefig(path_image, bbox_inches="tight", pad_inches=0.1)
|
||||
plt.close()
|
||||
|
||||
|
||||
async def example_multimodal_rag() -> None:
|
||||
"""Example for multimodal RAG"""
|
||||
# Reading the image and converting it to documents
|
||||
reader = ImageReader()
|
||||
docs = await reader(image_url=path_image)
|
||||
|
||||
# Create a knowledge base and add documents
|
||||
knowledge = SimpleKnowledge(
|
||||
embedding_model=DashScopeMultiModalEmbedding(
|
||||
api_key=os.environ["DASHSCOPE_API_KEY"],
|
||||
model_name="multimodal-embedding-v1",
|
||||
dimensions=1024,
|
||||
),
|
||||
embedding_store=QdrantStore(
|
||||
location=":memory:",
|
||||
collection_name="test_collection",
|
||||
dimensions=1024,
|
||||
),
|
||||
)
|
||||
|
||||
await knowledge.add_documents(docs)
|
||||
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You're a helpful assistant named Friday.",
|
||||
model=DashScopeChatModel(
|
||||
api_key=os.environ["DASHSCOPE_API_KEY"],
|
||||
model_name="qwen3-vl-plus",
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
knowledge=knowledge,
|
||||
)
|
||||
|
||||
await agent(
|
||||
Msg(
|
||||
"user",
|
||||
"Do you know my name?",
|
||||
"user",
|
||||
),
|
||||
)
|
||||
|
||||
# Let's see if the agent has stored the retrieved document in its memory
|
||||
print("\nThe retrieved document stored in the agent's memory:")
|
||||
content = (await agent.memory.get_memory())[-4].content
|
||||
print(json.dumps(content, indent=2, ensure_ascii=False))
|
||||
|
||||
|
||||
asyncio.run(example_multimodal_rag())
|
||||
78
examples/functionality/rag/react_agent_integration.py
Normal file
78
examples/functionality/rag/react_agent_integration.py
Normal file
@@ -0,0 +1,78 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The example of integrating ReAct agent with RAG."""
|
||||
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 SimpleKnowledge, QdrantStore, TextReader
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry point for the ReAct agent with RAG example."""
|
||||
|
||||
# Create an in-memory knowledge base instance
|
||||
print("Creating the knowledge base...")
|
||||
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",
|
||||
),
|
||||
)
|
||||
|
||||
# Insert some documents into the knowledge base
|
||||
# This could be done offline and only once
|
||||
print("Inserting documents into the knowledge base...")
|
||||
reader = TextReader(chunk_size=100, split_by="char")
|
||||
documents = await reader(
|
||||
# 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. ",
|
||||
)
|
||||
|
||||
print("Inserting documents into the knowledge base...")
|
||||
await knowledge.add_documents(documents)
|
||||
|
||||
# Integrate into the ReActAgent by the `knowledge` argument
|
||||
print("Creating the agent...")
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
model=DashScopeChatModel(
|
||||
api_key=os.environ["DASHSCOPE_API_KEY"],
|
||||
model_name="qwen-max",
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
# Equip the agent with the knowledge base
|
||||
knowledge=knowledge,
|
||||
print_hint_msg=True,
|
||||
)
|
||||
user = UserAgent(name="user")
|
||||
|
||||
# Start the conversation
|
||||
print("Start the conversation...")
|
||||
msg = Msg("user", "Do you know who is my best friend?", "user")
|
||||
while True:
|
||||
msg = await agent(msg)
|
||||
msg = await user(msg)
|
||||
if msg.get_text_content() == "exit":
|
||||
break
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
25
examples/functionality/session_with_sqlite/README.md
Normal file
25
examples/functionality/session_with_sqlite/README.md
Normal file
@@ -0,0 +1,25 @@
|
||||
# Session Management with Sqlite DB
|
||||
|
||||
This example demonstrates how to implement session management with a database backend. We use SQLite for simplicity,
|
||||
but the approach can be adapted for other databases.
|
||||
|
||||
Specifically, we implement a ``SqliteSession`` class that persists and retrieves session data from a SQLite table.
|
||||
The table schema includes fields for session ID, session data (stored as JSON), and timestamps for creation and last
|
||||
update.
|
||||
|
||||
We will create a simple agent and chat with it, then store the session data in the SQLite database. Then in the
|
||||
``test_load_session`` function, we will load the session data from the database and continue the chat.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Install agentscope from Pypi or source code.
|
||||
|
||||
```bash
|
||||
pip install agentscope
|
||||
```
|
||||
|
||||
Run the example by the following command
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
76
examples/functionality/session_with_sqlite/main.py
Normal file
76
examples/functionality/session_with_sqlite/main.py
Normal file
@@ -0,0 +1,76 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The main entry point for the session with SQLite example."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from sqlite_session import SqliteSession
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
|
||||
SQLITE_PATH = "./session.db"
|
||||
|
||||
|
||||
async def main(username: str, query: str) -> None:
|
||||
"""Create an agent, load from session, chat with it, and save its state
|
||||
to SQLite.
|
||||
|
||||
Args:
|
||||
username (`str`):
|
||||
The username to identify the session.
|
||||
query (`str`):
|
||||
The user input query.
|
||||
"""
|
||||
|
||||
agent = ReActAgent(
|
||||
name="friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
model=DashScopeChatModel(
|
||||
model_name="qwen-max",
|
||||
api_key=os.environ["DASHSCOPE_API_KEY"],
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
)
|
||||
|
||||
# Create the SQLite session
|
||||
session = SqliteSession(SQLITE_PATH)
|
||||
|
||||
# Load the agent state by the given key "friday_of_user"
|
||||
# The load_session_state supports multiple state modules
|
||||
await session.load_session_state(
|
||||
session_id=username,
|
||||
friday_of_user=agent,
|
||||
)
|
||||
|
||||
# Chat with it to generate some state
|
||||
await agent(
|
||||
Msg("user", query, "user"),
|
||||
)
|
||||
|
||||
# Save the agent state by the given key "friday_of_user"
|
||||
# Also support multiple state modules (e.g. multiple agents)
|
||||
await session.save_session_state(
|
||||
session_id=username,
|
||||
friday_of_user=agent,
|
||||
)
|
||||
|
||||
|
||||
print("User named Alice chats with the agent ...")
|
||||
asyncio.run(main("alice", "What's the capital of America?"))
|
||||
|
||||
print("User named Bob chats with the agent ...")
|
||||
asyncio.run(main("bob", "What's the capital of China?"))
|
||||
|
||||
print(
|
||||
"\nNow, let's recover the session for Alice and ask about what the user "
|
||||
"asked before.",
|
||||
)
|
||||
asyncio.run(
|
||||
main(
|
||||
"alice",
|
||||
"What did I ask you before, what's your answer and how many "
|
||||
"questions have I asked you?",
|
||||
),
|
||||
)
|
||||
167
examples/functionality/session_with_sqlite/sqlite_session.py
Normal file
167
examples/functionality/session_with_sqlite/sqlite_session.py
Normal file
@@ -0,0 +1,167 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The SQLite session class."""
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
|
||||
from agentscope import logger
|
||||
from agentscope.module import StateModule
|
||||
from agentscope.session import SessionBase
|
||||
|
||||
|
||||
class SqliteSession(SessionBase):
|
||||
"""A session that uses SQLite for storage."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sqlite_path: str,
|
||||
) -> None:
|
||||
"""Initialize the session.
|
||||
|
||||
Args:
|
||||
sqlite_path (`str`):
|
||||
The path to the SQLite database file.
|
||||
"""
|
||||
self.sqlite_path = sqlite_path
|
||||
|
||||
async def save_session_state(
|
||||
self,
|
||||
session_id: str,
|
||||
**state_modules_mapping: StateModule,
|
||||
) -> None:
|
||||
"""Save the session state to the SQLite database."""
|
||||
with sqlite3.connect(self.sqlite_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
# Prepare the session data as a dictionary
|
||||
session_data = {
|
||||
name: module.state_dict()
|
||||
for name, module in state_modules_mapping.items()
|
||||
}
|
||||
|
||||
json_data = json.dumps(session_data)
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS as_session (
|
||||
session_id TEXT,
|
||||
session_data JSON,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (session_id)
|
||||
)
|
||||
""",
|
||||
)
|
||||
|
||||
# Insert or replace the session data
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO as_session (session_id, session_data, updated_at)
|
||||
VALUES (?, json(?), CURRENT_TIMESTAMP)
|
||||
ON CONFLICT(session_id) DO UPDATE SET
|
||||
session_data = excluded.session_data,
|
||||
updated_at = excluded.updated_at
|
||||
""",
|
||||
(session_id, json_data),
|
||||
)
|
||||
conn.commit()
|
||||
cursor.close()
|
||||
|
||||
async def load_session_state(
|
||||
self,
|
||||
session_id: str,
|
||||
allow_not_exist: bool = True,
|
||||
**state_modules_mapping: StateModule,
|
||||
) -> None:
|
||||
"""Get the state dictionary from the SQLite database.
|
||||
|
||||
Args:
|
||||
session_id (`str`):
|
||||
The session id.
|
||||
allow_not_exist (`bool`, defaults to `True`):
|
||||
Whether to allow the session to not exist. If `False`, raises
|
||||
an error if the session does not exist.
|
||||
**state_modules_mapping (`list[StateModule]`):
|
||||
The list of state modules to be loaded.
|
||||
"""
|
||||
if not os.path.exists(self.sqlite_path):
|
||||
if allow_not_exist:
|
||||
logger.info(
|
||||
"SQLite database %s does not exist. "
|
||||
"Skipping load for session_id %s.",
|
||||
self.sqlite_path,
|
||||
session_id,
|
||||
)
|
||||
return
|
||||
raise ValueError(
|
||||
"Failed to load session state because the SQLite database "
|
||||
f"file '{self.sqlite_path}' does not exist.",
|
||||
)
|
||||
|
||||
with sqlite3.connect(self.sqlite_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
# If the table does not exist, return
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT name FROM sqlite_master WHERE type='table' AND
|
||||
name='as_session';
|
||||
""",
|
||||
)
|
||||
if cursor.fetchone() is None:
|
||||
if allow_not_exist:
|
||||
logger.info(
|
||||
"Session table does not exist in database %s. "
|
||||
"Skipping load for session_id %s.",
|
||||
self.sqlite_path,
|
||||
session_id,
|
||||
)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
"Failed to load session state because the session "
|
||||
"table 'as_session' does not exist in database "
|
||||
f"{self.sqlite_path}.",
|
||||
)
|
||||
|
||||
# Query the session data
|
||||
cursor.execute(
|
||||
"SELECT session_data FROM as_session WHERE session_id = ?",
|
||||
(session_id,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row is None:
|
||||
if allow_not_exist:
|
||||
logger.info(
|
||||
"Session_id %s does not exist in database %s. "
|
||||
"Skip loading.",
|
||||
session_id,
|
||||
self.sqlite_path,
|
||||
)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
f"Failed to load session state for session_id "
|
||||
f"{session_id} does not exist.",
|
||||
)
|
||||
|
||||
session_data = json.loads(row[0])
|
||||
|
||||
for name, module in state_modules_mapping.items():
|
||||
if name in session_data:
|
||||
module.load_state_dict(session_data[name])
|
||||
else:
|
||||
raise ValueError(
|
||||
f"State module '{name}' not found in session "
|
||||
"data.",
|
||||
)
|
||||
logger.info(
|
||||
"Load session state for session_id %s from "
|
||||
"database %s successfully.",
|
||||
session_id,
|
||||
self.sqlite_path,
|
||||
)
|
||||
|
||||
finally:
|
||||
cursor.close()
|
||||
@@ -0,0 +1,317 @@
|
||||
# MemoryWithCompress
|
||||
|
||||
- [ ] TODO: The memory module with compression will be added to the agentscope library in the future.
|
||||
|
||||
## Overview
|
||||
|
||||
MemoryWithCompress is a memory management system designed for AgentScope's `ReActAgent`. It automatically compresses conversation history when the memory size exceeds a specified token limit, using a Large Language Model (LLM) to create concise summaries that preserve key information. This allows agents to maintain context over long conversations while staying within token constraints.
|
||||
|
||||
The system maintains two separate storage mechanisms:
|
||||
- **`chat_history_storage`**: Stores the complete, unmodified conversation history (uses `MessageStorageBase` interface)
|
||||
- **`memory_storage`**: Stores messages that may be compressed when token limits are exceeded (uses `MessageStorageBase` interface)
|
||||
|
||||
Both storage mechanisms are abstracted through the `MessageStorageBase` interface, allowing for flexible storage backends. By default, `InMemoryMessageStorage` is used for both.
|
||||
|
||||
## Core Features
|
||||
|
||||
### Automatic Memory Compression
|
||||
- **Token-based Triggering**: Automatically compresses memory when the total token count exceeds `max_token`
|
||||
- **LLM-Powered Summarization**: Uses an LLM to intelligently compress conversation history while preserving essential information
|
||||
- **Structured Output**: Uses Pydantic schemas to ensure consistent compression format
|
||||
|
||||
### Dual Storage System
|
||||
- **Complete History**: Maintains original, unmodified messages in `_chat_history` for reference
|
||||
- **Compressed Memory**: Stores potentially compressed messages in `_memory` for efficient context management
|
||||
|
||||
### Flexible Memory Management
|
||||
- **Filtering Support**: Provides `filter_func` parameter for custom memory filtering
|
||||
- **Recent N Retrieval**: Supports retrieving only the most recent N messages
|
||||
- **State Persistence**: Includes `state_dict()` and `load_state_dict()` methods for saving and loading memory state
|
||||
- **Storage Abstraction**: Uses `MessageStorageBase` interface for flexible storage backends
|
||||
- **Compression Triggers**: Supports both token-based and custom trigger functions for compression
|
||||
- **Compression Timing Control**: Configurable compression on add (`compression_on_add`) and get (`compression_on_get`) operations
|
||||
|
||||
## File Structure
|
||||
|
||||
```
|
||||
memory_with_compression/
|
||||
├── README.md # This documentation file
|
||||
├── main.py # Example demonstrating MemoryWithCompress usage
|
||||
├── _memory_with_compress.py # Core MemoryWithCompress implementation
|
||||
├── _memory_storage.py # Storage abstraction layer (MessageStorageBase, InMemoryMessageStorage)
|
||||
├── _mc_utils.py # Utility functions (formatting, token counting, compression schema)
|
||||
|
||||
```
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### Clone the AgentScope Repository
|
||||
This example depends on AgentScope. Please clone the full repository to your local machine.
|
||||
|
||||
### Install Dependencies
|
||||
**Recommended**: Python 3.10+
|
||||
|
||||
Install the required dependencies:
|
||||
```bash
|
||||
pip install agentscope
|
||||
```
|
||||
|
||||
### API Keys
|
||||
This example uses DashScope APIs by default. You need to set your API key as an environment variable:
|
||||
```bash
|
||||
export DASHSCOPE_API_KEY='YOUR_API_KEY'
|
||||
```
|
||||
|
||||
You can easily switch to other models by modifying the configuration in `main.py`.
|
||||
|
||||
## How It Works
|
||||
|
||||
### 1. Memory Addition Flow
|
||||
1. **Message Input**: New messages are added via the async `add()` method
|
||||
2. **Dual Storage**: Messages are deep-copied and added to both `chat_history_storage` and `memory_storage`
|
||||
3. **Optional Compression on Add**: If `compression_on_add=True`, compression may be triggered immediately after adding messages
|
||||
|
||||
### 2. Memory Retrieval and Compression Flow
|
||||
When `get_memory()` is called (if `compression_on_get=True`):
|
||||
1. **Token Counting**: The system calculates the total token count of all messages in `memory_storage`
|
||||
2. **Compression Check**:
|
||||
- First checks if token count exceeds `max_token` (automatic compression)
|
||||
- Then checks if `compression_trigger_func` returns `True` (custom trigger)
|
||||
3. **LLM Compression**: If compression is needed, all messages in `memory_storage` are sent to the LLM with a compression prompt
|
||||
4. **Structured Output**: The LLM returns a structured response containing the compressed summary
|
||||
5. **Memory Replacement**: The entire `memory_storage` is updated with the compressed message(s)
|
||||
6. **Filtering & Selection**: Optional filtering and recent_n selection are applied
|
||||
7. **Return**: The processed memory is returned
|
||||
|
||||
### 3. Compression Process
|
||||
The compression uses a structured output approach:
|
||||
- **Prompt**: Instructs the LLM to summarize conversation history while preserving key information
|
||||
- **Customizable Prompt**: Supports `customized_compression_prompt` parameter for custom prompt templates
|
||||
- **Schema**: Uses `MemoryCompressionSchema` (Pydantic model) to ensure consistent output format
|
||||
- **Output Format**: Returns a message with content wrapped in `<compressed_memory>` tags
|
||||
- **Async Support**: All compression operations are asynchronous
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Running the Example
|
||||
To see `MemoryWithCompress` in action, run the example script:
|
||||
```bash
|
||||
python ./main.py
|
||||
```
|
||||
|
||||
### Basic Initialization
|
||||
Here is a snippet from `main.py` showing how to set up the agent and memory:
|
||||
|
||||
```python
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.token import OpenAITokenCounter
|
||||
from agentscope.message import Msg
|
||||
from _memory_with_compress import MemoryWithCompress
|
||||
|
||||
# 1. Create the model for agent and memory compression
|
||||
model = DashScopeChatModel(
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
model_name="qwen-max",
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# 2. Optional: Define a custom compression trigger function
|
||||
async def trigger_compression(msgs: list[Msg]) -> bool:
|
||||
# Trigger compression if the number of messages exceeds 2
|
||||
# and the last message is from the assistant
|
||||
return len(msgs) > 2 and msgs[-1].role == "assistant"
|
||||
|
||||
# 3. Initialize MemoryWithCompress
|
||||
memory_with_compress = MemoryWithCompress(
|
||||
model=model,
|
||||
formatter=DashScopeChatFormatter(),
|
||||
max_token=3000, # Compress when memory exceeds 3000 tokens
|
||||
token_counter=OpenAITokenCounter(model_name="qwen-max"),
|
||||
compression_trigger_func=trigger_compression, # Optional custom trigger
|
||||
compression_on_add=False, # Don't compress on add (default)
|
||||
compression_on_get=True, # Compress on get (default)
|
||||
)
|
||||
|
||||
# 4. Initialize ReActAgent with the memory instance
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
model=model,
|
||||
formatter=DashScopeChatFormatter(),
|
||||
memory=memory_with_compress,
|
||||
)
|
||||
```
|
||||
|
||||
### Custom Compression Function
|
||||
You can provide a custom compression function:
|
||||
|
||||
```python
|
||||
async def custom_compress(messages: List[Msg]) -> List[Msg]:
|
||||
# Your custom compression logic
|
||||
# Must return a List[Msg], not a single Msg
|
||||
compressed_content = "..."
|
||||
return [Msg("assistant", compressed_content, "assistant")]
|
||||
|
||||
memory_with_compress = MemoryWithCompress(
|
||||
model=model,
|
||||
formatter=formatter,
|
||||
max_token=300,
|
||||
compress_func=custom_compress,
|
||||
)
|
||||
```
|
||||
|
||||
### Custom Storage Backend
|
||||
You can provide custom storage backends by implementing the `MessageStorageBase` interface:
|
||||
|
||||
```python
|
||||
from _memory_storage import MessageStorageBase
|
||||
|
||||
class CustomStorage(MessageStorageBase):
|
||||
# Implement required methods: start, stop, health, add, delete, clear, get, replace, __aenter__, __aexit__
|
||||
...
|
||||
|
||||
memory_with_compress = MemoryWithCompress(
|
||||
model=model,
|
||||
formatter=formatter,
|
||||
max_token=300,
|
||||
chat_history_storage=CustomStorage(),
|
||||
memory_storage=CustomStorage(),
|
||||
)
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### MemoryWithCompress Class
|
||||
|
||||
#### `__init__(...)`
|
||||
Initializes the memory system. Key parameters include:
|
||||
|
||||
- `model` (ChatModelBase): The LLM model to use for compression
|
||||
- `formatter` (FormatterBase): The formatter to use for formatting messages
|
||||
- `max_token` (int): The maximum token count for `memory_storage`. Default: 28000. Compression is triggered when exceeded
|
||||
- `chat_history_storage` (MessageStorageBase): Storage backend for complete chat history. Default: `InMemoryMessageStorage()`
|
||||
- `memory_storage` (MessageStorageBase): Storage backend for compressed memory. Default: `InMemoryMessageStorage()`
|
||||
- `token_counter` (Optional[TokenCounterBase]): The token counter for counting tokens. Default: None. If None, it will return the character count of the JSON string representation of messages (i.e., len(json.dumps(messages, ensure_ascii=False))).
|
||||
- `compress_func` (Callable[[List[Msg]], Awaitable[List[Msg]]] | None): Custom compression function. Must be async and return `List[Msg]`. If None, uses the default `_compress_memory` method
|
||||
- `compression_trigger_func` (Callable[[List[Msg]], Awaitable[bool]] | None): Optional function to trigger compression when token count is below `max_token`. Must be async and return `bool`. If None, compression only occurs when token count exceeds `max_token`
|
||||
- `compression_on_add` (bool): Whether to check and compress memory when adding messages. Default: False
|
||||
- `compression_on_get` (bool): Whether to check and compress memory when getting messages. Default: True
|
||||
- `customized_compression_prompt` (str | None): Optional customized compression prompt template. Should include placeholders: `{max_token}`, `{messages_list_json}`, `{schema_json}`. Default: None (uses default template)
|
||||
|
||||
#### Main Methods
|
||||
|
||||
**`async add(msgs: Union[Sequence[Msg], Msg, None], compress_func=None, compression_trigger_func=None)`**
|
||||
- Adds new messages to both `chat_history_storage` and `memory_storage`
|
||||
- Messages are deep-copied to avoid modifying originals
|
||||
- Raises `TypeError` if non-Msg objects are provided
|
||||
- Parameters:
|
||||
- `msgs`: Messages to be added
|
||||
- `compress_func` (Optional): Override the instance-level compression function for this call
|
||||
- `compression_trigger_func` (Optional): Override the instance-level trigger function for this call
|
||||
- If `compression_on_add=True`, may trigger compression after adding
|
||||
|
||||
**`async direct_update_memory(msgs: Union[Sequence[Msg], Msg, None])`**
|
||||
- Directly updates the `memory_storage` with new messages (does not update `chat_history_storage`)
|
||||
- Useful for replacing memory content directly
|
||||
|
||||
**`async get_memory(recent_n=None, filter_func=None, compress_func=None, compression_trigger_func=None)`**
|
||||
- Retrieves memory content, automatically compressing if token limit is exceeded (if `compression_on_get=True`)
|
||||
- Parameters:
|
||||
- `recent_n` (Optional[int]): Return only the most recent N messages
|
||||
- `filter_func` (Optional[Callable[[int, Msg], bool]]): Custom filter function that takes (index, message) and returns bool
|
||||
- `compress_func` (Optional): Override the instance-level compression function for this call
|
||||
- `compression_trigger_func` (Optional): Override the instance-level trigger function for this call
|
||||
- Returns: `list[Msg]` - The memory content (potentially compressed)
|
||||
|
||||
**`async delete(indices: Union[Iterable[int], int])`**
|
||||
- Deletes memory fragments from `memory_storage` (note: does not delete from `chat_history_storage`)
|
||||
- Indices can be a single int or an iterable of ints
|
||||
|
||||
**`async size() -> int`**
|
||||
- Returns the number of messages in `chat_history_storage`
|
||||
|
||||
**`async clear()`**
|
||||
- Clears all memory from both `chat_history_storage` and `memory_storage`
|
||||
|
||||
**`state_dict() -> dict`**
|
||||
- Returns a dictionary containing the serialized state:
|
||||
- `chat_history_storage`: List of message dictionaries from chat history
|
||||
- `memory_storage`: List of message dictionaries from memory
|
||||
- `max_token`: The max_token setting
|
||||
- Note: This method handles async operations internally, so it can be called from both sync and async contexts
|
||||
|
||||
**`load_state_dict(state_dict: dict, strict: bool = True)`**
|
||||
- Loads memory state from a dictionary
|
||||
- Restores `chat_history_storage`, `memory_storage`, and `max_token` settings
|
||||
- Note: This method handles async operations internally, so it can be called from both sync and async contexts
|
||||
|
||||
**`async retrieve(*args, **kwargs)`**
|
||||
- Not implemented. Use `get_memory()` instead.
|
||||
- Raises `NotImplementedError`
|
||||
|
||||
## Internal Methods
|
||||
|
||||
**`async _compress_memory(msgs: List[Msg]) -> List[Msg]`**
|
||||
- Internal method that compresses messages using the LLM
|
||||
- Uses structured output with `MemoryCompressionSchema`
|
||||
- Returns a `List[Msg]` containing the compressed summary (typically a single message)
|
||||
- Supports both streaming and non-streaming models
|
||||
|
||||
**`async _check_length_and_compress(compress_func=None) -> bool`**
|
||||
- Checks if memory token count exceeds `max_token` and compresses if needed
|
||||
- Returns `True` if compression was triggered, `False` otherwise
|
||||
|
||||
**`async check_and_compress(compress_func=None, compression_trigger_func=None, memory=None) -> tuple[bool, List[Msg]]`**
|
||||
- Checks if compression is needed based on `compression_trigger_func`
|
||||
- Returns a tuple: (was_compressed: bool, compressed_memory: List[Msg])
|
||||
- If `memory` is provided, checks that instead of `memory_storage`
|
||||
|
||||
## Utility Functions
|
||||
|
||||
The `_mc_utils.py` module provides:
|
||||
|
||||
- **`format_msgs(msgs)`**: Formats a list of `Msg` objects into a list of dictionaries
|
||||
- **`async count_words(token_counter, text)`**: Counts tokens in text (supports both string and list[dict] formats). Must be awaited.
|
||||
- **`MemoryCompressionSchema`**: Pydantic model for structured compression output
|
||||
- **`DEFAULT_COMPRESSION_PROMPT_TEMPLATE`**: Default prompt template for compression (includes placeholders: `{max_token}`, `{messages_list_json}`, `{schema_json}`)
|
||||
|
||||
## Storage Abstraction
|
||||
|
||||
The `_memory_storage.py` module provides:
|
||||
|
||||
- **`MessageStorageBase`**: Abstract base class for message storage backends
|
||||
- Required async methods: `start()`, `stop()`, `health()`, `add()`, `delete()`, `clear()`, `get()`, `replace()`, `__aenter__()`, `__aexit__()`
|
||||
- **`InMemoryMessageStorage`**: Default in-memory implementation
|
||||
- Stores messages in a simple list
|
||||
- Suitable for most use cases
|
||||
|
||||
## Best Practices
|
||||
|
||||
- **Token Limit Selection**: Choose `max_token` based on your model's context window and typical conversation length
|
||||
- **Compression Timing**:
|
||||
- Set `compression_on_get=True` (default) for compression during retrieval
|
||||
- Set `compression_on_add=False` (default) to avoid compression during add operations, as it may not complete before `get_memory()` is called
|
||||
- **Async Operations**: All main methods are async, so use `await` when calling them
|
||||
- **State Persistence**: Use `state_dict()` and `load_state_dict()` to save/restore conversation state between sessions
|
||||
- **Custom Compression**: For domain-specific compression needs, implement a custom `compress_func` (must be async and return `List[Msg]`)
|
||||
- **Compression Triggers**: Use `compression_trigger_func` for custom compression logic beyond token limits (e.g., compress after N messages, compress on specific conditions)
|
||||
- **Storage Backends**: Implement custom `MessageStorageBase` subclasses for persistent storage (e.g., database, file system)
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- **Compression Not Triggering**:
|
||||
- Check that `compression_on_get=True` if you expect compression during retrieval
|
||||
- Verify that `max_token` is set appropriately
|
||||
- Ensure `get_memory()` is being called (and awaited)
|
||||
- If using `compression_trigger_func`, verify it returns `True` when compression should occur
|
||||
- **Structured Output Errors**: Ensure your model supports structured output (e.g., DashScope models with `structured_model` parameter)
|
||||
- **Token Counting Issues**: Verify that your `token_counter` is compatible with your model and correctly configured
|
||||
- **Async/Await Errors**: Remember that most methods are async - use `await` when calling them
|
||||
- **Storage Issues**: If using custom storage backends, ensure all required methods are properly implemented and async
|
||||
|
||||
## Reference
|
||||
|
||||
- [AgentScope Documentation](https://github.com/agentscope-ai/agentscope)
|
||||
- [Pydantic Documentation](https://docs.pydantic.dev/)
|
||||
@@ -0,0 +1,46 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The main entry point of the MemoryWithCompress example."""
|
||||
import asyncio
|
||||
import os
|
||||
from agentscope.agent import ReActAgent, UserAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.token import CharTokenCounter
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry point of the MemoryWithCompress example."""
|
||||
|
||||
# Create model for agent and memory compression
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
model=DashScopeChatModel(
|
||||
api_key=os.getenv("DASHSCOPE_API_KEY"),
|
||||
model_name="qwen3-max",
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
compression_config=ReActAgent.CompressionConfig(
|
||||
enable=True,
|
||||
agent_token_counter=CharTokenCounter(),
|
||||
# We set a small trigger threshold for demonstration purposes.
|
||||
trigger_threshold=1000,
|
||||
keep_recent=3,
|
||||
),
|
||||
)
|
||||
user = UserAgent("User")
|
||||
|
||||
# Simulate a conversation to trigger memory compression
|
||||
msg = None
|
||||
while True:
|
||||
msg = await user(msg)
|
||||
if msg.get_text_content() == "exit":
|
||||
break
|
||||
msg = await agent(msg)
|
||||
|
||||
print("The memory of the agent:")
|
||||
for msg in await agent.memory.get_memory():
|
||||
print(msg.to_dict(), end="\n")
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
479
examples/functionality/short_term_memory/reme/README.md
Normal file
479
examples/functionality/short_term_memory/reme/README.md
Normal file
@@ -0,0 +1,479 @@
|
||||
# ReMe Short-Term Memory in AgentScope
|
||||
|
||||
This example demonstrates how to
|
||||
|
||||
- use ReMeShortTermMemory to provide automatic working memory management for AgentScope agents,
|
||||
- handle long conversation contexts with intelligent summarization and compaction,
|
||||
- integrate short-term memory with ReAct agents for efficient tool usage and context management, and
|
||||
- configure DashScope models for memory operations.
|
||||
|
||||
## Why Short-Term Memory?
|
||||
|
||||
### The Challenge: From Prompt Engineering to Context Engineering
|
||||
|
||||
As AI agents evolved from simple chatbots to sophisticated autonomous systems, the focus shifted from "prompt engineering" to "context engineering". While prompt engineering focused on crafting effective instructions for language models, context engineering addresses a more fundamental challenge: **managing the ever-growing conversation and tool execution history that agents accumulate**.
|
||||
|
||||
### The Core Problem: Context Explosion
|
||||
|
||||
Agentic systems work by binding LLMs with tools and running them in a loop where the agent decides which tools to call and feeds results back into the message history. This creates a snowball effect:
|
||||
|
||||
- **Rapid Growth**: A seemingly simple task can trigger 50+ tool calls, with production agents often running hundreds of conversation turns
|
||||
- **Large Outputs**: Each tool call can return substantial text, consuming massive amounts of tokens
|
||||
- **Memory Pressure**: The context window quickly fills up as messages and tool results accumulate chronologically
|
||||
|
||||
### The Consequence: Context Rot
|
||||
|
||||
When context grows too large, model performance degrades significantly—a phenomenon known as **"context rot"**:
|
||||
|
||||
- **Repetitive Responses**: The model starts generating redundant or circular answers
|
||||
- **Slower Reasoning**: Inference becomes noticeably slower as context length increases
|
||||
- **Quality Degradation**: Overall response quality and coherence decline
|
||||
- **Lost Focus**: The model struggles to identify relevant information in the bloated context
|
||||
|
||||
### The Fundamental Paradox
|
||||
|
||||
Agents face a critical tension:
|
||||
|
||||
- **Need Rich Context**: Agents require comprehensive historical information to make informed decisions
|
||||
- **Suffer from Large Context**: Excessive context causes performance degradation and inefficiency
|
||||
|
||||
**Context management aims to keep "just enough" information in the window**—sufficient for effective decision-making while leaving room for retrieval and expansion, without overwhelming the model.
|
||||
|
||||
### Why Short-Term Memory Management Matters
|
||||
|
||||
Effective short-term memory management is essential for:
|
||||
|
||||
1. **Maintaining Performance**: Keeping context within optimal size prevents quality degradation
|
||||
2. **Enabling Long-Running Tasks**: Agents can handle complex, multi-step workflows without hitting context limits
|
||||
3. **Cost Efficiency**: Reducing token usage directly lowers API costs
|
||||
4. **Preserving Reasoning Quality**: Clean, focused context helps models maintain coherent reasoning chains
|
||||
5. **Scalability**: Proper memory management allows agents to scale to production workloads
|
||||
|
||||
### The Solution: Intelligent Context Management
|
||||
|
||||
ReMeShortTermMemory implements proven context management strategies:
|
||||
|
||||
- **Context Offloading**: Moving large tool outputs to external storage while keeping references
|
||||
- **Context Reduction**: Compacting tool results into minimal representations and summarizing when necessary
|
||||
- **Smart Retention**: Keeping recent messages intact to maintain continuity and provide usage examples
|
||||
- **Automatic Triggering**: Monitoring token usage and applying strategies before performance degrades
|
||||
|
||||
By implementing these strategies, ReMeShortTermMemory enables agents to handle arbitrarily long conversations and complex tasks while maintaining optimal performance throughout.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Python 3.10 or higher
|
||||
- DashScope API key from Alibaba Cloud
|
||||
|
||||
|
||||
## QuickStart
|
||||
|
||||
Install agentscope and ensure you have a valid DashScope API key in your environment variables.
|
||||
|
||||
> Note: The example is built with DashScope chat model. If you want to use OpenAI models instead,
|
||||
> modify the model initialization in the example code accordingly.
|
||||
|
||||
```bash
|
||||
# Install agentscope from source
|
||||
cd {PATH_TO_AGENTSCOPE}
|
||||
pip install -e .
|
||||
# Install dependencies
|
||||
pip install reme-ai python-dotenv
|
||||
```
|
||||
|
||||
Set up your API key:
|
||||
|
||||
```bash
|
||||
export DASHSCOPE_API_KEY='YOUR_API_KEY'
|
||||
```
|
||||
|
||||
Or create a `.env` file:
|
||||
|
||||
```bash
|
||||
DASHSCOPE_API_KEY=YOUR_API_KEY
|
||||
```
|
||||
|
||||
Run the example:
|
||||
|
||||
```bash
|
||||
python short_term_memory_example.py
|
||||
```
|
||||
|
||||
The example will:
|
||||
1. Initialize a ReMeShortTermMemory instance with DashScope models
|
||||
2. Demonstrate automatic memory compaction for long tool responses
|
||||
3. Show integration with ReActAgent for context-aware conversations
|
||||
4. Use grep and read_file tools to search and retrieve information from files
|
||||
|
||||
## Key Features
|
||||
|
||||
- **Automatic Memory Management**: Intelligent summarization and compaction of working memory to handle long contexts
|
||||
- **Tool Response Optimization**: Automatic truncation and summarization of large tool responses to stay within token limits
|
||||
- **Flexible Configuration**: Configurable thresholds for compaction ratio, token limits, and recent message retention
|
||||
- **ReAct Agent Integration**: Seamless integration with AgentScope's ReActAgent and tool system
|
||||
- **Async Operations**: Full async support for non-blocking memory operations
|
||||
|
||||
## Basic Usage
|
||||
|
||||
This section provides a detailed walkthrough of the `short_term_memory_example.py` code, explaining how each component works together to create an agent with intelligent context management.
|
||||
|
||||
### Configuration Parameters
|
||||
|
||||
#### `ReMeShortTermMemory` Class Parameters
|
||||
|
||||
The `ReMeShortTermMemory` class accepts the following initialization parameters:
|
||||
|
||||
- **`model`** (`DashScopeChatModel | OpenAIChatModel | None`): Language model for compression operations. Must be either `DashScopeChatModel` or `OpenAIChatModel`. This model is used for LLM-based compression when generating compact state snapshots. **Required**.
|
||||
|
||||
- **`reme_config_path`** (`str | None`): Optional path to ReMe configuration file for custom settings. Use this to provide advanced ReMe configurations beyond the standard parameters. Default: `None`.
|
||||
|
||||
- **`working_summary_mode`** (`str`): Strategy for working memory management. Controls how the memory system handles context overflow:
|
||||
- `"compact"`: Only compact verbose tool messages by storing full content externally and keeping short previews in the active context.
|
||||
- `"compress"`: Only apply LLM-based compression to generate compact state snapshots of conversation history.
|
||||
- `"auto"`: First run compaction, then optionally run compression if the compaction ratio exceeds `compact_ratio_threshold`. This is the recommended mode for most use cases.
|
||||
|
||||
Default: `"auto"`.
|
||||
|
||||
- **`compact_ratio_threshold`** (`float`): Threshold for compaction effectiveness in AUTO mode. If `(compacted_tokens / original_tokens) > compact_ratio_threshold`, compression is applied after compaction. This ensures compression only runs when compaction alone isn't sufficient. Valid range: 0.0 to 1.0. Default: `0.75`.
|
||||
|
||||
- **`max_total_tokens`** (`int`): Maximum token count threshold before compression is triggered. This limit does **not** include `keep_recent_count` messages or system messages, which are always preserved. Should be set to 20%-50% of your model's context window size to leave room for new tool calls and responses. Default: `20000`.
|
||||
|
||||
- **`max_tool_message_tokens`** (`int`): Maximum token count for individual tool messages before compaction. Tool messages exceeding this limit are stored externally in files, with only a short preview kept in the active context. This is the maximum tolerable length for a single tool response. Default: `2000`.
|
||||
|
||||
- **`group_token_threshold`** (`int | None`): Maximum token count per compression group when splitting messages for LLM compression. When set to a positive integer, long message sequences are split into smaller batches for compression. If `None` or `0`, all messages are compressed in a single group. Use this to control the granularity of compression operations. Default: `None`.
|
||||
|
||||
- **`keep_recent_count`** (`int`): Number of most recent messages to preserve without compression or compaction. These messages remain in full in the active context to maintain conversation continuity and provide usage examples for the agent. The example uses `1` for demonstration purposes; **in production, a value of `10` is recommended** to maintain better conversation flow. Default: `10`.
|
||||
|
||||
- **`store_dir`** (`str`): Directory path for storing offloaded message content and compressed history files. This is where external files containing full tool responses and compressed message history are saved. The directory will be created automatically if it doesn't exist. Default: `"inmemory"`.
|
||||
|
||||
- **`**kwargs`** (`Any`): Additional arguments passed to `ReMeApp` initialization. Use this to pass any extra configuration options supported by the underlying ReMe application.
|
||||
|
||||
#### Parameter Relationships and Best Practices
|
||||
|
||||
- **Token Budget Strategy**: Set `max_total_tokens` to 20%-50% of your model's context window. For example, if your model has a 128K context window, set `max_total_tokens` between 25,600 and 64,000 tokens.
|
||||
|
||||
- **Compaction vs Compression**:
|
||||
- Compaction is fast and lossless (full content is stored externally)
|
||||
- Compression is slower but more aggressive (uses LLM to summarize)
|
||||
- Use `"auto"` mode to benefit from both strategies
|
||||
|
||||
- **Recent Message Retention**: Higher `keep_recent_count` values (e.g., 10) provide better context continuity but consume more tokens. Lower values (e.g., 1) are more aggressive but may lose important recent context.
|
||||
|
||||
- **Tool Message Handling**: Adjust `max_tool_message_tokens` based on your typical tool response sizes. If your tools frequently return large outputs (e.g., file contents, API responses), consider a higher threshold or ensure compaction is enabled.
|
||||
|
||||
### Code Flow Diagram
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A[Start: Load Environment] --> B[Create Toolkit]
|
||||
B --> C[Register Tools: grep & read_file]
|
||||
C --> D[Initialize LLM Model]
|
||||
D --> E[Create ReMeShortTermMemory]
|
||||
E --> F[Enter Async Context Manager]
|
||||
F --> G[Add Initial Messages with Large Tool Response]
|
||||
G --> H[Memory Auto-Compacts Large Content]
|
||||
H --> I[Create ReActAgent with Memory]
|
||||
I --> J[User Sends Query]
|
||||
J --> K[Agent Uses Tools to Search/Read]
|
||||
K --> L[Tool Responses Added to Memory]
|
||||
L --> M{Memory Token Limit?}
|
||||
M -->|Exceeded| N[Auto-Compact/Summarize]
|
||||
M -->|OK| O[Agent Generates Response]
|
||||
N --> O
|
||||
O --> P[Return Response to User]
|
||||
P --> Q[Exit Context Manager]
|
||||
Q --> End[End]
|
||||
|
||||
style H fill:#e1f5ff
|
||||
style N fill:#ffe1e1
|
||||
style O fill:#e1ffe1
|
||||
```
|
||||
|
||||
### Step-by-Step Code Walkthrough
|
||||
|
||||
The example demonstrates a complete workflow from tool registration to agent interaction. Here's a detailed breakdown:
|
||||
|
||||
#### 1. Environment Setup and Imports
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
```
|
||||
|
||||
The code starts by loading environment variables (including the DashScope API key) from a `.env` file.
|
||||
|
||||
#### 2. Tool Registration
|
||||
|
||||
The example defines two custom tools that demonstrate how to integrate retrieval operations:
|
||||
|
||||
**`grep` Tool**: Searches for patterns in files using regular expressions
|
||||
```python
|
||||
async def grep(file_path: str, pattern: str, limit: str) -> ToolResponse:
|
||||
"""A powerful search tool for finding patterns in files..."""
|
||||
from reme_ai.retrieve.working import GrepOp
|
||||
|
||||
op = GrepOp()
|
||||
await op.async_call(file_path=file_path, pattern=pattern, limit=limit)
|
||||
return ToolResponse(
|
||||
content=[TextBlock(type="text", text=op.output)],
|
||||
)
|
||||
```
|
||||
|
||||
**`read_file` Tool**: Reads specific line ranges from files
|
||||
```python
|
||||
async def read_file(file_path: str, offset: int, limit: int) -> ToolResponse:
|
||||
"""Reads and returns the content of a specified file..."""
|
||||
from reme_ai.retrieve.working import ReadFileOp
|
||||
|
||||
op = ReadFileOp()
|
||||
await op.async_call(file_path=file_path, offset=offset, limit=limit)
|
||||
return ToolResponse(
|
||||
content=[TextBlock(type="text", text=op.output)],
|
||||
)
|
||||
```
|
||||
|
||||
> **Important Note on Tool Replaceability**:
|
||||
> - The `grep` and `read_file` tools shown here are **example implementations** using ReMe's built-in operations
|
||||
> - You can **replace them with your own retrieval tools**, such as:
|
||||
> - Vector database embedding retrieval (e.g., ChromaDB, Pinecone, Weaviate)
|
||||
> - Web search APIs (e.g., Google Search, Bing Search)
|
||||
> - Database query tools (e.g., SQL queries, MongoDB queries)
|
||||
> - Custom domain-specific search solutions
|
||||
> - Similarly, the **offline write operations** (used internally by ReMeShortTermMemory to store compacted content) can be customized by modifying the `write_text_file` function in AgentScope's tool system
|
||||
> - The key requirement is that your tools return `ToolResponse` objects with appropriate content blocks
|
||||
|
||||
#### 3. LLM Model Initialization
|
||||
|
||||
```python
|
||||
llm = DashScopeChatModel(
|
||||
model_name="qwen3-coder-30b-a3b-instruct",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
generate_kwargs={
|
||||
"temperature": 0.001,
|
||||
"seed": 0,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
The model is configured with low temperature for consistent, deterministic responses. This same model will be used for both agent reasoning and memory summarization operations.
|
||||
|
||||
#### 4. Short-Term Memory Initialization
|
||||
|
||||
```python
|
||||
short_term_memory = ReMeShortTermMemory(
|
||||
model=llm,
|
||||
working_summary_mode="auto", # Automatic memory management
|
||||
compact_ratio_threshold=0.75, # Trigger compaction at 75% capacity
|
||||
max_total_tokens=20000, # Set to 20%-50% of model's context window
|
||||
max_tool_message_tokens=2000, # Maximum tolerable tool response length
|
||||
group_token_threshold=None, # Max tokens per LLM compression batch; None means no splitting
|
||||
keep_recent_count=1, # Keep 1 recent message intact (set to 1 for demo; use 10 in production)
|
||||
store_dir="inmemory", # Storage directory for offloaded content
|
||||
)
|
||||
```
|
||||
|
||||
This configuration enables automatic memory management that will:
|
||||
- Monitor token usage
|
||||
- Automatically compact large tool responses when they exceed `max_tool_message_tokens`
|
||||
- Trigger summarization when total tokens exceed `max_total_tokens` and compaction ratio exceeds `compact_ratio_threshold`
|
||||
|
||||
#### 5. Async Context Manager Usage
|
||||
|
||||
```python
|
||||
async with short_term_memory:
|
||||
# All memory operations happen here
|
||||
```
|
||||
|
||||
The `async with` statement ensures proper initialization and cleanup of memory resources. This is the recommended approach for using `ReMeShortTermMemory`.
|
||||
|
||||
#### 6. Simulating Long Context
|
||||
|
||||
The example demonstrates memory compaction by adding a large tool response:
|
||||
|
||||
```python
|
||||
# Read README content and multiply it 10 times to simulate a large response
|
||||
f = open("../../../../README.md", encoding="utf-8")
|
||||
readme_content = f.read()
|
||||
f.close()
|
||||
|
||||
memories = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "搜索下项目资料",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [...], # Tool call metadata
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": readme_content * 10, # Large tool response (10x README)
|
||||
"tool_call_id": "call_6596dafa2a6a46f7a217da",
|
||||
},
|
||||
]
|
||||
|
||||
await short_term_memory.add(
|
||||
ReMeShortTermMemory.list_to_msg(memories),
|
||||
allow_duplicates=True,
|
||||
)
|
||||
```
|
||||
|
||||
When this large content is added, `ReMeShortTermMemory` will:
|
||||
1. Detect that the tool response exceeds `max_tool_message_tokens` (the maximum tolerable tool response length, set to 2000 in this example)
|
||||
2. Automatically compact it by storing the full content in an external file
|
||||
3. Keep only a short preview in the active memory
|
||||
4. This happens transparently without manual intervention
|
||||
|
||||
#### 7. ReAct Agent Creation
|
||||
|
||||
```python
|
||||
agent = ReActAgent(
|
||||
name="react",
|
||||
sys_prompt=(
|
||||
"You are a helpful assistant. "
|
||||
"工具调用的调用可能会被缓存到本地。"
|
||||
"可以先使用`Grep`匹配关键词或者正则表达式所在行数，然后通过`ReadFile`读取位置附近的代码。"
|
||||
# ... more instructions
|
||||
),
|
||||
model=llm,
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=short_term_memory, # Memory is integrated here
|
||||
max_iters=20,
|
||||
)
|
||||
```
|
||||
|
||||
The agent is configured with:
|
||||
- The same LLM model used for memory operations
|
||||
- The toolkit containing `grep` and `read_file` tools
|
||||
- The `short_term_memory` instance for automatic context management
|
||||
- A system prompt that guides the agent on tool usage patterns
|
||||
|
||||
#### 8. Agent Interaction
|
||||
|
||||
```python
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content=("项目资料中，agentscope_v1论文的一作是谁？"),
|
||||
name="user",
|
||||
)
|
||||
msg = await agent(msg)
|
||||
print(f"✓ Agent response: {msg.get_text_content()}\n")
|
||||
```
|
||||
|
||||
When the agent processes this message:
|
||||
1. It receives the user query
|
||||
2. Decides to use tools (e.g., `grep` to search for "agentscope_v1")
|
||||
3. Tool responses are automatically added to memory
|
||||
4. If memory grows too large, automatic compaction occurs
|
||||
5. The agent generates a response based on the managed context
|
||||
6. The response is returned to the user
|
||||
|
||||
### Complete Example Code Structure
|
||||
|
||||
```python
|
||||
async def main() -> None:
|
||||
# 1. Create toolkit and register tools
|
||||
toolkit = Toolkit()
|
||||
toolkit.register_tool_function(grep)
|
||||
toolkit.register_tool_function(read_file)
|
||||
|
||||
# 2. Initialize LLM
|
||||
llm = DashScopeChatModel(...)
|
||||
|
||||
# 3. Create short-term memory
|
||||
short_term_memory = ReMeShortTermMemory(...)
|
||||
|
||||
# 4. Use async context manager
|
||||
async with short_term_memory:
|
||||
# 5. Add initial messages (with large content to demo compaction)
|
||||
await short_term_memory.add(messages, allow_duplicates=True)
|
||||
|
||||
# 6. Create agent with memory
|
||||
agent = ReActAgent(..., memory=short_term_memory, ...)
|
||||
|
||||
# 7. Interact with agent
|
||||
response = await agent(user_message)
|
||||
```
|
||||
|
||||
### Key Takeaways
|
||||
|
||||
1. **Automatic Memory Management**: Memory compaction and summarization happen automatically when thresholds are exceeded
|
||||
2. **Tool Integration**: Tools return `ToolResponse` objects that are seamlessly integrated into memory
|
||||
3. **Async Context Manager**: Always use `async with short_term_memory:` to ensure proper resource management
|
||||
4. **Flexible Tool System**: The `grep` and `read_file` tools are examples—you can replace them with any retrieval mechanism that fits your use case
|
||||
5. **Transparent Operation**: Memory management is transparent to the agent—it just sees a clean, focused context
|
||||
|
||||
### Using Async Context Manager
|
||||
|
||||
`ReMeShortTermMemory` implements the async context manager protocol, which ensures proper initialization and cleanup of resources. There are two ways to use it:
|
||||
|
||||
#### Recommended: Using `async with` Statement
|
||||
|
||||
The recommended approach is to use the `async with` statement, which automatically handles resource management:
|
||||
|
||||
```python
|
||||
async with short_term_memory:
|
||||
# Memory is initialized here
|
||||
await short_term_memory.add(messages)
|
||||
response = await agent(msg)
|
||||
# Memory is automatically cleaned up when exiting the block
|
||||
```
|
||||
|
||||
#### Alternative: Manual `__aenter__` and `__aexit__` Calls
|
||||
|
||||
You can also manually call `__aenter__` and `__aexit__` if you need more control:
|
||||
|
||||
```python
|
||||
# Manually initialize
|
||||
await short_term_memory.__aenter__()
|
||||
|
||||
try:
|
||||
# Use the memory
|
||||
await short_term_memory.add(messages)
|
||||
response = await agent(msg)
|
||||
finally:
|
||||
# Manually cleanup
|
||||
await short_term_memory.__aexit__(None, None, None)
|
||||
```
|
||||
|
||||
> **Note**: It's recommended to use the `async with` statement as it ensures proper resource cleanup even if an exception occurs.
|
||||
|
||||
## Advanced Configuration
|
||||
|
||||
You can customize the ReMe config by passing a config path:
|
||||
|
||||
```python
|
||||
short_term_memory = ReMeShortTermMemory(
|
||||
model=llm,
|
||||
reme_config_path="path/to/your/config.yaml", # Pass your custom ReMe configuration
|
||||
# ... other parameters
|
||||
)
|
||||
```
|
||||
|
||||
For more configuration options, refer to the [ReMe documentation](https://github.com/agentscope-ai/ReMe).
|
||||
|
||||
## What's in the Example
|
||||
|
||||
The `short_term_memory_example.py` file demonstrates:
|
||||
|
||||
1. **Tool Integration**: Registering `grep` and `read_file` tools for searching and reading files
|
||||
2. **Memory Initialization**: Setting up ReMeShortTermMemory with appropriate parameters for handling long contexts
|
||||
3. **Long Context Handling**: Adding a large tool response (README content × 10) to demonstrate automatic memory compaction
|
||||
4. **ReAct Agent Usage**: Using the agent with short-term memory to answer questions based on retrieved information
|
||||
|
||||
## Example Workflow
|
||||
|
||||
The example shows a typical workflow:
|
||||
|
||||
1. User asks to search for project information
|
||||
2. Agent uses `grep` tool to find relevant content
|
||||
3. Agent uses `read_file` tool to read specific sections
|
||||
4. Large tool responses are automatically compacted by the memory system
|
||||
5. Agent answers the user's question based on the retrieved information
|
||||
|
||||
@@ -0,0 +1,349 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""ReMe-based short-term memory implementation for AgentScope."""
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, List
|
||||
from uuid import uuid4
|
||||
|
||||
from agentscope import logger
|
||||
from agentscope._utils._common import _json_loads_with_repair
|
||||
from agentscope.formatter import DashScopeChatFormatter, OpenAIChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.message import Msg, TextBlock, ToolUseBlock, ToolResultBlock
|
||||
from agentscope.model import DashScopeChatModel, OpenAIChatModel
|
||||
from agentscope.tool import write_text_file
|
||||
|
||||
|
||||
class ReMeShortTermMemory(InMemoryMemory):
|
||||
"""Short-term memory implementation using ReMe for message management.
|
||||
|
||||
This class provides automatic working-memory management through a
|
||||
multi-stage pipeline that reduces token usage while preserving
|
||||
essential information:
|
||||
|
||||
1. **Compaction**: Truncates large tool messages by storing full
|
||||
content in external files and keeping only short previews in the
|
||||
active context.
|
||||
2. **Compression**: Uses LLM to generate dense summaries of older
|
||||
conversation history, creating a compact state snapshot.
|
||||
3. **Offload**: Orchestrates compaction and optional compression
|
||||
based on the configured working_summary_mode (COMPACT, COMPRESS,
|
||||
or AUTO).
|
||||
|
||||
The memory management is triggered automatically when `get_memory()`
|
||||
is called, ensuring the agent's context stays within token limits
|
||||
while maintaining access to detailed historical information through
|
||||
external storage.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: DashScopeChatModel | OpenAIChatModel | None = None,
|
||||
reme_config_path: str | None = None,
|
||||
working_summary_mode: str = "auto",
|
||||
compact_ratio_threshold: float = 0.75,
|
||||
max_total_tokens: int = 20000,
|
||||
max_tool_message_tokens: int = 2000,
|
||||
group_token_threshold: int | None = None,
|
||||
keep_recent_count: int = 10,
|
||||
store_dir: str = "inmemory",
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize ReMe-based short-term memory.
|
||||
|
||||
Args:
|
||||
model: Language model for compression operations. Must be
|
||||
either DashScopeChatModel or OpenAIChatModel.
|
||||
reme_config_path: Optional path to ReMe configuration file
|
||||
for custom settings.
|
||||
working_summary_mode: Strategy for working memory management.
|
||||
- "compact": Only compact verbose tool messages by
|
||||
storing full content externally and keeping short
|
||||
previews.
|
||||
- "compress": Only apply LLM-based compression to
|
||||
generate compact state snapshots.
|
||||
- "auto": First run compaction, then optionally run
|
||||
compression if the compaction ratio exceeds
|
||||
compact_ratio_threshold.
|
||||
Defaults to "auto".
|
||||
compact_ratio_threshold: Threshold for compaction
|
||||
effectiveness in AUTO mode. If (compacted_tokens /
|
||||
original_tokens) > this threshold, compression is
|
||||
applied. Defaults to 0.75.
|
||||
max_total_tokens: Maximum token count threshold before
|
||||
compression is triggered. Does not include
|
||||
keep_recent_count messages or system messages.
|
||||
Defaults to 20000.
|
||||
max_tool_message_tokens: Maximum token count for individual
|
||||
tool messages before compaction. Tool messages exceeding
|
||||
this are stored externally. Defaults to 2000.
|
||||
group_token_threshold: Maximum token count per compression
|
||||
group when splitting messages for LLM compression. If
|
||||
None or 0, all messages are compressed in a single
|
||||
group. Defaults to None.
|
||||
keep_recent_count: Number of most recent messages to
|
||||
preserve without compression or compaction. These
|
||||
messages remain in full in the active context.
|
||||
Defaults to 1.
|
||||
store_dir: Directory path for storing offloaded message
|
||||
content and compressed history files. Defaults to
|
||||
"working_memory".
|
||||
**kwargs: Additional arguments passed to ReMeApp
|
||||
initialization.
|
||||
|
||||
Raises:
|
||||
ValueError: If model is not a DashScopeChatModel or
|
||||
OpenAIChatModel.
|
||||
ImportError: If reme_ai library is not installed.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Store working memory parameters
|
||||
self.working_summary_mode = working_summary_mode
|
||||
self.compact_ratio_threshold = compact_ratio_threshold
|
||||
self.max_total_tokens = max_total_tokens
|
||||
self.max_tool_message_tokens = max_tool_message_tokens
|
||||
self.group_token_threshold = group_token_threshold
|
||||
self.keep_recent_count = keep_recent_count
|
||||
self.store_dir = store_dir
|
||||
|
||||
config_args = []
|
||||
|
||||
if isinstance(model, DashScopeChatModel):
|
||||
llm_api_base = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
llm_api_key = model.api_key
|
||||
self.formatter = DashScopeChatFormatter()
|
||||
|
||||
elif isinstance(model, OpenAIChatModel):
|
||||
llm_api_base = str(getattr(model.client, "base_url", None))
|
||||
llm_api_key = str(getattr(model.client, "api_key", None))
|
||||
self.formatter = OpenAIChatFormatter()
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
"model must be a DashScopeChatModel or "
|
||||
"OpenAIChatModel instance. "
|
||||
f"Got {type(model).__name__} instead.",
|
||||
)
|
||||
|
||||
llm_model_name = model.model_name
|
||||
|
||||
if llm_model_name:
|
||||
config_args.append(f"llm.default.model_name={llm_model_name}")
|
||||
|
||||
try:
|
||||
from reme_ai import ReMeApp
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"The 'reme_ai' library is required for ReMe-based "
|
||||
"short-term memory. Please try `pip install reme-ai`,"
|
||||
"and visit: https://github.com/agentscope-ai/ReMe for more "
|
||||
"information.",
|
||||
) from e
|
||||
|
||||
self.app = ReMeApp(
|
||||
*config_args,
|
||||
llm_api_key=llm_api_key,
|
||||
llm_api_base=llm_api_base,
|
||||
embedding_api_key=llm_api_key, # fake api key
|
||||
embedding_api_base=llm_api_base, # fake api base
|
||||
config_path=reme_config_path,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self._app_started = False
|
||||
|
||||
async def __aenter__(self) -> "ReMeShortTermMemory":
|
||||
"""Async context manager entry.
|
||||
|
||||
Initializes the ReMe application for async operations.
|
||||
"""
|
||||
if self.app is not None:
|
||||
await self.app.__aenter__()
|
||||
self._app_started = True
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: Any = None,
|
||||
exc_val: Any = None,
|
||||
exc_tb: Any = None,
|
||||
) -> None:
|
||||
"""Async context manager exit.
|
||||
|
||||
Cleans up the ReMe application resources.
|
||||
"""
|
||||
if self.app is not None:
|
||||
await self.app.__aexit__(exc_type, exc_val, exc_tb)
|
||||
self._app_started = False
|
||||
|
||||
async def get_memory(self) -> list[Msg]:
|
||||
"""Retrieve and manage working memory with automatic summarization.
|
||||
|
||||
This method performs the core working-memory management pipeline:
|
||||
|
||||
1. **Format messages**: Converts internal Msg objects to standard
|
||||
message format using the appropriate formatter (DashScope or
|
||||
OpenAI).
|
||||
2. **Execute offload pipeline**: Calls ReMe's
|
||||
summary_working_memory_for_as operation which orchestrates:
|
||||
- Message compaction: Large tool messages are truncated and
|
||||
stored externally with only previews kept in context.
|
||||
- Message compression: If needed (based on
|
||||
working_summary_mode), older messages are compressed using
|
||||
LLM into dense summaries.
|
||||
- File storage: Offloaded content is written to external
|
||||
files for potential retrieval.
|
||||
3. **Update content**: Replaces the internal message list with
|
||||
the managed version, ensuring subsequent operations work with
|
||||
the optimized context.
|
||||
|
||||
The operation respects configuration parameters like
|
||||
max_total_tokens, keep_recent_count, and working_summary_mode to
|
||||
balance context size with information preservation.
|
||||
|
||||
Returns:
|
||||
List of Msg objects representing the managed working memory,
|
||||
with large tool messages compacted and/or older history
|
||||
compressed as needed.
|
||||
|
||||
Note:
|
||||
This method automatically writes offloaded content to files
|
||||
in the configured store_dir. The write_file_dict metadata
|
||||
contains paths and content for all externally stored
|
||||
messages.
|
||||
"""
|
||||
messages: list[dict[str, Any]] = await self.formatter.format(
|
||||
msgs=self.content, # type: ignore[has-type]
|
||||
)
|
||||
for message in messages:
|
||||
if isinstance(message.get("content"), list):
|
||||
msg_content = message.get("content")
|
||||
logger.warning(
|
||||
"Skipping message with content as list. content=%s",
|
||||
msg_content,
|
||||
)
|
||||
message["content"] = ""
|
||||
|
||||
# Execute ReMe's working memory offload pipeline
|
||||
# This orchestrates compaction and/or compression based on
|
||||
# working_summary_mode
|
||||
result: dict = await self.app.async_execute(
|
||||
name="summary_working_memory_for_as",
|
||||
messages=messages,
|
||||
working_summary_mode=self.working_summary_mode,
|
||||
compact_ratio_threshold=self.compact_ratio_threshold,
|
||||
max_total_tokens=self.max_total_tokens,
|
||||
max_tool_message_tokens=self.max_tool_message_tokens,
|
||||
group_token_threshold=self.group_token_threshold,
|
||||
keep_recent_count=self.keep_recent_count,
|
||||
store_dir=self.store_dir,
|
||||
chat_id=uuid4().hex,
|
||||
)
|
||||
logger.info(
|
||||
"summary_working_memory_for_as.result=%s",
|
||||
json.dumps(result, ensure_ascii=False, indent=2),
|
||||
)
|
||||
|
||||
# Extract managed messages and file write operations from result
|
||||
messages = result.get("answer", [])
|
||||
write_file_dict: dict = result.get("metadata", {}).get(
|
||||
"write_file_dict",
|
||||
{},
|
||||
)
|
||||
# Write offloaded content to external files
|
||||
# This includes full tool message content and compressed message
|
||||
# history
|
||||
if write_file_dict:
|
||||
for path, content_str in write_file_dict.items():
|
||||
file_dir = Path(path).parent
|
||||
if not file_dir.exists():
|
||||
file_dir.mkdir(parents=True, exist_ok=True)
|
||||
await write_text_file(path, content_str)
|
||||
|
||||
# Update internal content with managed messages
|
||||
self.content = self.list_to_msg(messages)
|
||||
return self.content
|
||||
|
||||
@staticmethod
|
||||
def list_to_msg(messages: list[dict[str, Any]]) -> list[Msg]:
|
||||
"""Convert a list of message dictionaries to Msg objects.
|
||||
|
||||
This method handles the conversion from standard message format
|
||||
(used by ReMe and LLM APIs) back to AgentScope's Msg objects.
|
||||
It properly handles:
|
||||
- Text content for user, system, and assistant messages
|
||||
- Tool result blocks (converting role="tool" to role="system")
|
||||
- Tool use blocks from tool_calls in assistant messages
|
||||
|
||||
Args:
|
||||
messages: List of message dictionaries with role, content,
|
||||
and optional tool_calls or tool-related fields.
|
||||
|
||||
Returns:
|
||||
List of Msg objects with properly structured content blocks.
|
||||
"""
|
||||
msg_list: list[Msg] = []
|
||||
for msg_dict in messages:
|
||||
role = msg_dict["role"]
|
||||
content_blocks: List[
|
||||
TextBlock | ToolUseBlock | ToolResultBlock
|
||||
] = []
|
||||
content = msg_dict.get("content")
|
||||
|
||||
# Convert text content to appropriate content blocks
|
||||
if content:
|
||||
if role in ["user", "system", "assistant"]:
|
||||
content_blocks.append(TextBlock(type="text", text=content))
|
||||
elif role in ["tool"]:
|
||||
# Tool messages are converted to system messages with
|
||||
# ToolResultBlock
|
||||
role = "system"
|
||||
content_blocks.append(
|
||||
ToolResultBlock(
|
||||
type="tool_result",
|
||||
name=msg_dict.get("name"),
|
||||
id=msg_dict.get("tool_call_id"),
|
||||
output=[TextBlock(type="text", text=content)],
|
||||
),
|
||||
)
|
||||
|
||||
# Convert tool_calls to ToolUseBlock content blocks
|
||||
if msg_dict.get("tool_calls"):
|
||||
for tool_call in msg_dict["tool_calls"]:
|
||||
# Parse tool arguments with repair for malformed JSON
|
||||
input_ = _json_loads_with_repair(
|
||||
tool_call["function"].get(
|
||||
"arguments",
|
||||
"{}",
|
||||
)
|
||||
or "{}",
|
||||
)
|
||||
content_blocks.append(
|
||||
ToolUseBlock(
|
||||
type="tool_use",
|
||||
name=tool_call["function"]["name"],
|
||||
input=input_,
|
||||
id=tool_call["id"],
|
||||
),
|
||||
)
|
||||
|
||||
msg_obj = Msg(
|
||||
name=role,
|
||||
content=content_blocks,
|
||||
role=role,
|
||||
metadata=msg_dict.get("metadata"),
|
||||
)
|
||||
msg_list.append(msg_obj)
|
||||
return msg_list
|
||||
|
||||
async def retrieve(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Retrieve operation is not implemented for ReMe short-term memory.
|
||||
|
||||
ReMe focuses on working memory management (compaction and compression)
|
||||
rather than retrieval from long-term storage.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: This operation is not supported.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,188 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Example demonstrating ReMeShortTermMemory usage with ReActAgent."""
|
||||
# noqa: E402
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.message import Msg, TextBlock
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import ToolResponse, Toolkit, view_text_file
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Main function demonstrating ReMeShortTermMemory with tool usage."""
|
||||
from reme_short_term_memory import ReMeShortTermMemory
|
||||
|
||||
toolkit = Toolkit()
|
||||
|
||||
async def grep(file_path: str, pattern: str, limit: str) -> ToolResponse:
|
||||
"""A powerful search tool for finding patterns in files using regular
|
||||
expressions.
|
||||
|
||||
Supports full regex syntax (e.g., "log.*Error", "function\\s+\\w+"),
|
||||
glob pattern filtering, and result limiting. Ideal for searching code
|
||||
or text content across multiple files.
|
||||
|
||||
Args:
|
||||
file_path (`str`):
|
||||
The path to the file to search in. Can be an absolute or
|
||||
relative path.
|
||||
pattern (`str`):
|
||||
The search pattern or regular expression to match. Supports
|
||||
full regex syntax for complex pattern matching.
|
||||
limit (`str`):
|
||||
The maximum number of matching results to return. Use this to
|
||||
control output size for large files. Should not exceed 50.
|
||||
"""
|
||||
from reme_ai.retrieve.working import GrepOp
|
||||
|
||||
op = GrepOp()
|
||||
await op.async_call(file_path=file_path, pattern=pattern, limit=limit)
|
||||
return ToolResponse(
|
||||
content=[
|
||||
TextBlock(
|
||||
type="text",
|
||||
text=op.output,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
async def read_file(
|
||||
file_path: str,
|
||||
offset: int,
|
||||
limit: int,
|
||||
) -> ToolResponse:
|
||||
"""Reads and returns the content of a specified file.
|
||||
|
||||
For text files, it can read specific line ranges using the 'offset' and
|
||||
'limit' parameters. Use offset and limit to paginate through large
|
||||
files.
|
||||
|
||||
Note: It's recommended to use the `grep` tool first to locate the line
|
||||
numbers of interest before calling this function.
|
||||
|
||||
Args:
|
||||
file_path (`str`):
|
||||
The path to the file to read. Can be an absolute or relative
|
||||
path.
|
||||
offset (`int`):
|
||||
The starting line number to read from (0-indexed). Use this to
|
||||
skip to a specific position in the file.
|
||||
limit (`int`):
|
||||
The maximum number of lines to read from the offset position.
|
||||
Helps control memory usage when reading large files. Should
|
||||
not exceed 100.
|
||||
"""
|
||||
|
||||
return await view_text_file(file_path, ranges=[offset, offset + limit])
|
||||
|
||||
# These two tools are provided as examples. You can replace them with your
|
||||
# own retrieval tools, such as vector database embedding retrieval or other
|
||||
# search solutions that fit your use case.
|
||||
toolkit.register_tool_function(grep)
|
||||
toolkit.register_tool_function(read_file)
|
||||
|
||||
llm = DashScopeChatModel(
|
||||
model_name="qwen3-max",
|
||||
# model_name="qwen3-coder-30b-a3b-instruct",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
stream=False,
|
||||
generate_kwargs={
|
||||
"temperature": 0.001,
|
||||
"seed": 0,
|
||||
},
|
||||
)
|
||||
short_term_memory = ReMeShortTermMemory(
|
||||
model=llm,
|
||||
working_summary_mode="auto",
|
||||
compact_ratio_threshold=0.75,
|
||||
max_total_tokens=20000,
|
||||
max_tool_message_tokens=2000,
|
||||
group_token_threshold=None, # Max tokens per compression batch
|
||||
keep_recent_count=1, # Set to 1 for demo; use 10 in production
|
||||
store_dir="inmemory",
|
||||
)
|
||||
|
||||
async with short_term_memory:
|
||||
# Simulate ultra long context
|
||||
f = open("../../../../README.md", encoding="utf-8")
|
||||
readme_content = f.read()
|
||||
f.close()
|
||||
|
||||
memories = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Search for project information",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_6596dafa2a6a46f7a217da",
|
||||
"function": {
|
||||
"arguments": "{}",
|
||||
"name": "web_search",
|
||||
},
|
||||
"type": "function",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": readme_content * 10,
|
||||
"tool_call_id": "call_6596dafa2a6a46f7a217da",
|
||||
},
|
||||
]
|
||||
await short_term_memory.add(
|
||||
ReMeShortTermMemory.list_to_msg(memories),
|
||||
allow_duplicates=True,
|
||||
)
|
||||
|
||||
agent = ReActAgent(
|
||||
name="react",
|
||||
sys_prompt=(
|
||||
"You are a helpful assistant. "
|
||||
"Tool calls may be cached locally. "
|
||||
"You can first use `Grep` to match keywords or regular "
|
||||
"expressions to find line numbers, then use `ReadFile` "
|
||||
"to read the code near that location. "
|
||||
"If no matches are found, never give up trying - try "
|
||||
"other parameters or relax the matching conditions, such "
|
||||
"as searching for only partial keywords. "
|
||||
"After `Grep`, you can use the `ReadFile` command to "
|
||||
"view content starting from a specified offset position "
|
||||
"`offset` with length `limit`. "
|
||||
"The maximum limit is 100. "
|
||||
"If the current content is insufficient, the `ReadFile` "
|
||||
"command can continuously try different `offset` and "
|
||||
"`limit` parameters."
|
||||
),
|
||||
model=llm,
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=short_term_memory,
|
||||
max_iters=20,
|
||||
)
|
||||
|
||||
msg = Msg(
|
||||
role="user",
|
||||
content=(
|
||||
"In the project documentation, who is the first author "
|
||||
"of the agentscope_v1 paper?"
|
||||
),
|
||||
name="user",
|
||||
)
|
||||
msg = await agent(msg)
|
||||
print(f"✓ Agent response: {msg.get_text_content()}\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
28
examples/functionality/stream_printing_messages/README.md
Normal file
28
examples/functionality/stream_printing_messages/README.md
Normal file
@@ -0,0 +1,28 @@
|
||||
# Stream Printing Messages
|
||||
|
||||
The AgentScope agent is designed to communicate with the user and the other agents by passing messages explicitly.
|
||||
However, we notice the requirements that obtain the printing messages from the agent in a streaming manner.
|
||||
Therefore, in example we demonstrate how to gather and yield the printing messages from a single agent and
|
||||
multi-agent systems in a streaming manner.
|
||||
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run the following command to see the streaming printing messages from the agent.
|
||||
Note the messages with the same ID are the chunks of the same message in accumulated manner.
|
||||
|
||||
- For single-agent:
|
||||
|
||||
```bash
|
||||
python single_agent.py
|
||||
```
|
||||
|
||||
- For multi-agent:
|
||||
|
||||
```bash
|
||||
python multi_agent.py
|
||||
```
|
||||
|
||||
> Note: The example is built with DashScope chat model. If you want to change the model in this example, don't forget
|
||||
> to change the formatter at the same time! The corresponding relationship between built-in models and formatters are
|
||||
> list in [our tutorial](https://doc.agentscope.io/tutorial/task_prompt.html#id1)
|
||||
@@ -0,0 +1,62 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Example for gather the printing messages from multiple agents."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeMultiAgentFormatter
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.pipeline import MsgHub, stream_printing_messages
|
||||
|
||||
|
||||
def create_agent(name: str) -> ReActAgent:
|
||||
"""Create an agent with the given name."""
|
||||
return ReActAgent(
|
||||
name=name,
|
||||
sys_prompt=f"You are a student named {name}.",
|
||||
model=DashScopeChatModel(
|
||||
api_key=os.environ["DASHSCOPE_API_KEY"],
|
||||
model_name="qwen-max",
|
||||
stream=False, # close streaming for simplicity
|
||||
),
|
||||
formatter=DashScopeMultiAgentFormatter(),
|
||||
)
|
||||
|
||||
|
||||
async def workflow(
|
||||
alice: ReActAgent,
|
||||
bob: ReActAgent,
|
||||
charlie: ReActAgent,
|
||||
) -> None:
|
||||
"""The example workflow for multiple agents."""
|
||||
async with MsgHub(
|
||||
participants=[alice, bob, charlie],
|
||||
announcement=Msg(
|
||||
"user",
|
||||
"Alice, Bob and Charlie, welcome to the meeting! Let's "
|
||||
"meet each other first.",
|
||||
"user",
|
||||
),
|
||||
):
|
||||
# agent speaks in turn
|
||||
await alice()
|
||||
await bob()
|
||||
await charlie()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry for the example."""
|
||||
# Create agents
|
||||
alice, bob, charlie = [
|
||||
create_agent(_) for _ in ["Alice", "Bob", "Charlie"]
|
||||
]
|
||||
|
||||
async for msg, last in stream_printing_messages(
|
||||
agents=[alice, bob, charlie],
|
||||
coroutine_task=workflow(alice, bob, charlie),
|
||||
):
|
||||
print(msg, last)
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,62 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The example demonstrating how to obtain the messages from the agent in a
|
||||
streaming way."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.pipeline import stream_printing_messages
|
||||
from agentscope.tool import (
|
||||
Toolkit,
|
||||
execute_shell_command,
|
||||
view_text_file,
|
||||
execute_python_code,
|
||||
)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main function."""
|
||||
toolkit = Toolkit()
|
||||
toolkit.register_tool_function(execute_shell_command)
|
||||
toolkit.register_tool_function(execute_python_code)
|
||||
toolkit.register_tool_function(view_text_file)
|
||||
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
# Change the model and formatter together if you want to try other
|
||||
# models
|
||||
model=DashScopeChatModel(
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
model_name="qwen-max",
|
||||
enable_thinking=False,
|
||||
stream=True,
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
)
|
||||
|
||||
# Prepare a user message
|
||||
user_msg = Msg(
|
||||
"user",
|
||||
"Hi! Who are you?",
|
||||
"user",
|
||||
)
|
||||
|
||||
# We disable the terminal printing to avoid messy outputs
|
||||
agent.set_console_output_enabled(False)
|
||||
|
||||
# obtain the printing messages from the agent in a streaming way
|
||||
async for msg, last in stream_printing_messages(
|
||||
agents=[agent],
|
||||
coroutine_task=agent(user_msg),
|
||||
):
|
||||
print(msg, last)
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
96
examples/functionality/structured_output/README.md
Normal file
96
examples/functionality/structured_output/README.md
Normal file
@@ -0,0 +1,96 @@
|
||||
# Structured Output Example
|
||||
|
||||
## What This Example Demonstrates
|
||||
|
||||
This example showcases **structured output generation** using AgentScope with Pydantic models. It demonstrates how to constrain AI model outputs to follow specific data structures and formats, ensuring consistent and parseable responses.
|
||||
|
||||
### Key Features:
|
||||
- **Structured Data Generation**: Forces agent responses to conform to
|
||||
predefined schemas
|
||||
- **Pydantic Integration**: Uses Pydantic models to define output structure with validation
|
||||
- **Type Safety**: Ensures output data types match expected formats
|
||||
- **Field Validation**: Includes constraints like age limits (0-120) and enum choices
|
||||
- **JSON Output**: Generates clean, structured JSON responses
|
||||
|
||||
### Example Models:
|
||||
|
||||
1. **TableModel**: Structured person information
|
||||
- `name`: Person's name (string)
|
||||
- `age`: Person's age (integer,0-120)
|
||||
- `intro`: One-sentence introduction (string)
|
||||
- `honors`: List of honors/achievements (array of strings)
|
||||
|
||||
2. **ChoiceModel**: Constrained choice selection
|
||||
- `choice`: Must be one of "apple", "banana", or "orange"
|
||||
|
||||
### Use Cases:
|
||||
- **Data Extraction**: Extract structured information from unstructured text
|
||||
- **Form Generation**: Generate consistent data for databases or APIs
|
||||
- **Survey Responses**: Ensure responses fit predefined categories
|
||||
- **Content Classification**: Categorize content into specific types
|
||||
|
||||
## How to Run This Example
|
||||
1. **Set Environment Variable:**
|
||||
```bash
|
||||
export DASHSCOPE_API_KEY="your_dashscope_api_key_here"
|
||||
```
|
||||
2. **Run the script:**
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
3. **Expected Output:**
|
||||
The program will generate two structured responses like below:
|
||||
```
|
||||
Structured Output 1:
|
||||
{
|
||||
"name": "Albert Einstein",
|
||||
"age": 76,
|
||||
"intro": 1,
|
||||
"honors": [
|
||||
"Nobel Prize in Physics (1921)",
|
||||
"Copley Medal (1925)"
|
||||
]
|
||||
}
|
||||
Structured Output 2:
|
||||
{
|
||||
"choice": "apple"
|
||||
}
|
||||
```
|
||||
|
||||
>💡**Note:** The specific content will vary with each run since the agent generates different responses, but the JSON structure will always conform to the predefined Pydantic models (`TableModel` and `ChoiceModel`).
|
||||
|
||||
## How It Works:
|
||||
1. The agent receives a query along with a structured_model parameter
|
||||
2. The agent generates a response that conforms to the Pydantic model schema
|
||||
3. The structured data is returned in res.metadata as a validated JSON object
|
||||
4. Pydantic ensures all field types and constraints are satisfied
|
||||
|
||||
## Custom Pydantic Models
|
||||
Create your own structured output models for specific use cases, for example:
|
||||
|
||||
```
|
||||
from typing import Optional
|
||||
from pydantic import BaseModel, Field, EmailStr
|
||||
|
||||
class BusinessModel(BaseModel):
|
||||
"""Business information extraction model."""
|
||||
|
||||
company_name: str = Field(description="Name of the company")
|
||||
industry: str = Field(description="Industry sector")
|
||||
founded_year: int = Field(description="Year founded", ge=1800, le=2024)
|
||||
headquarters: str = Field(description="Location of headquarters")
|
||||
employee_count: Optional[int] = Field(description="Number of employees", ge=1)
|
||||
email: Optional[EmailStr] = Field(description="Contact email address")
|
||||
website: Optional[str] = Field(description="Company website URL")
|
||||
|
||||
# Usage
|
||||
query = Msg("user", "Tell me about Tesla Inc.", "user")
|
||||
res = await agent(query, structured_model=BusinessModel)
|
||||
```
|
||||
|
||||
## Best Practices for Structured Output
|
||||
|
||||
1. **Use Descriptive Field Names:** Make field purposes clear
|
||||
2. **Add Field Descriptions:** Help the agent understand what data to generate
|
||||
3. **Set Validation Constraints:** Use Pydantic validators for data integrity
|
||||
4. **Choose Appropriate Types:** Use specific types like EmailStr, datetime, etc.
|
||||
80
examples/functionality/structured_output/main.py
Normal file
80
examples/functionality/structured_output/main.py
Normal file
@@ -0,0 +1,80 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The main entry point of the structured output example."""
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from agentscope.agent import ReActAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.message import Msg
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import Toolkit
|
||||
|
||||
|
||||
class TableModel(BaseModel):
|
||||
"""A simple table model for structured output."""
|
||||
|
||||
name: str = Field(description="The name of the person")
|
||||
age: int = Field(description="The age of the person", ge=0, le=120)
|
||||
intro: str = Field(description="A one-sentence introduction of the person")
|
||||
honors: list[str] = Field(
|
||||
description="A list of honors received by this person",
|
||||
)
|
||||
|
||||
|
||||
class ChoiceModel(BaseModel):
|
||||
"""A simple choice model for structured output."""
|
||||
|
||||
choice: Literal["apple", "banana", "orange"] = Field(
|
||||
description="Your choice of fruit",
|
||||
)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry point for the structured output example."""
|
||||
toolkit = Toolkit()
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
model=DashScopeChatModel(
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
model_name="qwen-max",
|
||||
stream=True,
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
)
|
||||
|
||||
query_msg_1 = Msg(
|
||||
"user",
|
||||
"Please introduce Einstein",
|
||||
"user",
|
||||
)
|
||||
res = await agent(query_msg_1, structured_model=TableModel)
|
||||
print(
|
||||
"Structured Output 1:\n"
|
||||
"```\n"
|
||||
f"{json.dumps(res.metadata, indent=4)}\n"
|
||||
"```",
|
||||
)
|
||||
|
||||
query_msg_2 = Msg(
|
||||
"user",
|
||||
"Choose one of your favorite fruit",
|
||||
"user",
|
||||
)
|
||||
res = await agent(query_msg_2, structured_model=ChoiceModel)
|
||||
print(
|
||||
"Structured Output 2:\n"
|
||||
"```\n"
|
||||
f"{json.dumps(res.metadata, indent=4)}\n"
|
||||
"```",
|
||||
)
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
13
examples/functionality/tts/README.md
Normal file
13
examples/functionality/tts/README.md
Normal file
@@ -0,0 +1,13 @@
|
||||
# TTS (Text-to-Speech) in AgentScope
|
||||
|
||||
This example demonstrates how to integrate DashScope Realtime TTS model with `ReActAgent` to enable audio output.
|
||||
The agent can speak its responses in real-time.
|
||||
|
||||
This example uses DashScope's Realtime TTS model, you can also change to other TTS models supported by AgentScope, e.g.
|
||||
OpenAI, Gemini, etc.
|
||||
|
||||
To run the example, execute:
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
56
examples/functionality/tts/main.py
Normal file
56
examples/functionality/tts/main.py
Normal file
@@ -0,0 +1,56 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""The main entry point of the ReAct agent example."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.agent import ReActAgent, UserAgent
|
||||
from agentscope.formatter import DashScopeChatFormatter
|
||||
from agentscope.memory import InMemoryMemory
|
||||
from agentscope.model import DashScopeChatModel
|
||||
from agentscope.tool import (
|
||||
Toolkit,
|
||||
execute_shell_command,
|
||||
execute_python_code,
|
||||
view_text_file,
|
||||
)
|
||||
from agentscope.tts import DashScopeRealtimeTTSModel
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""The main entry point for the ReAct agent example."""
|
||||
toolkit = Toolkit()
|
||||
toolkit.register_tool_function(execute_shell_command)
|
||||
toolkit.register_tool_function(execute_python_code)
|
||||
toolkit.register_tool_function(view_text_file)
|
||||
|
||||
agent = ReActAgent(
|
||||
name="Friday",
|
||||
sys_prompt="You are a helpful assistant named Friday.",
|
||||
model=DashScopeChatModel(
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
model_name="qwen3-max",
|
||||
enable_thinking=False,
|
||||
stream=True,
|
||||
),
|
||||
formatter=DashScopeChatFormatter(),
|
||||
toolkit=toolkit,
|
||||
memory=InMemoryMemory(),
|
||||
# Specify the TTS model for real-time speech synthesis
|
||||
tts_model=DashScopeRealtimeTTSModel(
|
||||
model_name="qwen3-tts-flash-realtime",
|
||||
api_key=os.environ.get("DASHSCOPE_API_KEY"),
|
||||
voice="Cherry",
|
||||
stream=False,
|
||||
),
|
||||
)
|
||||
user = UserAgent("User")
|
||||
|
||||
msg = None
|
||||
while True:
|
||||
msg = await user(msg)
|
||||
if msg.get_text_content() == "exit":
|
||||
break
|
||||
msg = await agent(msg)
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,422 @@
|
||||
# AlibabaCloud MySQL Vector Store Example
|
||||
|
||||
This example demonstrates how to use the `AlibabaCloudMySQLStore` class in AgentScope's RAG system for vector storage and similarity search operations using AlibabaCloud MySQL (RDS) with native vector functions.
|
||||
|
||||
## Features
|
||||
|
||||
AlibabaCloudMySQLStore provides:
|
||||
- Vector storage using MySQL's native VECTOR data type
|
||||
- Automatic vector index creation (CREATE VECTOR INDEX) based on distance metric
|
||||
- Vector functions (VEC_FROMTEXT, VEC_DISTANCE_COSINE, VEC_DISTANCE_EUCLIDEAN)
|
||||
- Database-level distance calculation and sorting via ORDER BY
|
||||
- Two distance metrics: COSINE and EUCLIDEAN (supported by AlibabaCloud MySQL)
|
||||
- Metadata filtering support
|
||||
- CRUD operations (Create, Read, Update, Delete)
|
||||
- Support for chunked documents
|
||||
- Direct access to the underlying MySQL connection for advanced operations
|
||||
- Full integration with AlibabaCloud RDS MySQL instances
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### 1. AlibabaCloud RDS MySQL Instance
|
||||
|
||||
You need an AlibabaCloud RDS MySQL instance with vector support:
|
||||
|
||||
- **Version**: MySQL 8.0+
|
||||
- **Vector Plugin**: Ensure the vector search plugin is enabled (check `vidx_disabled` parameter is OFF)
|
||||
- **Network Access**: Configure security group and whitelist to allow access
|
||||
|
||||
#### Create RDS MySQL Instance on AlibabaCloud:
|
||||
|
||||
1. Go to [AlibabaCloud RDS Console](https://rdsnext.console.aliyun.com/)
|
||||
2. Click "Create Instance"
|
||||
3. Select MySQL 8.0 or higher
|
||||
4. Configure specifications based on your needs
|
||||
5. Set up network and security settings
|
||||
6. Note down the connection endpoint (e.g., `rm-xxxxx.mysql.rds.aliyuncs.com`)
|
||||
|
||||
#### Configure Database:
|
||||
|
||||
```sql
|
||||
-- Connect to your RDS MySQL instance
|
||||
mysql -h rm-xxxxx.mysql.rds.aliyuncs.com -P 3306 -u your_username -p
|
||||
|
||||
-- Check if vector capability is enabled (vidx_disabled should be OFF)
|
||||
SHOW VARIABLES LIKE 'vidx_disabled';
|
||||
-- Expected result: vidx_disabled | OFF
|
||||
-- If OFF, vector capability is enabled
|
||||
-- If ON, contact AlibabaCloud support to enable vector search plugin
|
||||
|
||||
-- Create database
|
||||
CREATE DATABASE agentscope_test;
|
||||
|
||||
-- Use the database
|
||||
USE agentscope_test;
|
||||
|
||||
-- Verify vector functions are available
|
||||
SELECT VEC_FROMTEXT('[1,2,3]');
|
||||
```
|
||||
|
||||
### 2. Python Dependencies
|
||||
|
||||
```bash
|
||||
pip install mysql-connector-python agentscope
|
||||
```
|
||||
|
||||
### 3. Network Configuration
|
||||
|
||||
Ensure your local machine or server can access the RDS instance:
|
||||
- Add your IP to the RDS whitelist
|
||||
- Configure security group rules
|
||||
- Use SSL connection if required
|
||||
|
||||
## Configuration
|
||||
|
||||
Update the connection parameters in `main.py`:
|
||||
|
||||
```python
|
||||
store = AlibabaCloudMySQLStore(
|
||||
host="rm-xxxxx.mysql.rds.aliyuncs.com", # Your RDS endpoint
|
||||
port=3306,
|
||||
user="your_username", # Your RDS username
|
||||
password="your_password", # Your RDS password
|
||||
database="agentscope_test",
|
||||
table_name="test_vectors",
|
||||
dimensions=768, # Set to your embedding dimension
|
||||
distance="COSINE",
|
||||
# Optional: SSL configuration
|
||||
# connection_kwargs={
|
||||
# "ssl_ca": "/path/to/ca.pem",
|
||||
# "ssl_verify_cert": True,
|
||||
# }
|
||||
)
|
||||
```
|
||||
|
||||
## Running the Example
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
|
||||
## Example Tests
|
||||
|
||||
The example includes three comprehensive tests:
|
||||
|
||||
### 1. Basic CRUD Operations
|
||||
- Initialize AlibabaCloudMySQLStore
|
||||
- Add documents with embeddings
|
||||
- Search for similar documents
|
||||
- Delete documents
|
||||
- Get the underlying MySQL connection
|
||||
|
||||
### 2. Search with Metadata Filtering
|
||||
- Add documents with different categories
|
||||
- Search with and without filters
|
||||
- Use SQL WHERE clauses for filtering
|
||||
|
||||
### 3. Different Distance Metrics
|
||||
- Test COSINE similarity (best for normalized vectors)
|
||||
- Test EUCLIDEAN distance (best for absolute distance)
|
||||
|
||||
## Key Features Explained
|
||||
|
||||
### Distance Metrics
|
||||
|
||||
AlibabaCloud MySQL supports two distance metrics:
|
||||
|
||||
- **COSINE**: Measures the cosine of the angle between vectors. Values range from 0 (identical) to 2 (opposite). Best for text embeddings and normalized vectors.
|
||||
- **EUCLIDEAN**: Measures the straight-line Euclidean distance between vectors. Lower values indicate similarity. Best for absolute distance measurements.
|
||||
|
||||
**Note**: Unlike some other vector databases, AlibabaCloud MySQL currently only supports COSINE and EUCLIDEAN distance functions. Inner Product (IP) is not supported.
|
||||
|
||||
### Metadata Filtering
|
||||
|
||||
Use SQL WHERE clauses to filter search results:
|
||||
|
||||
```python
|
||||
results = await store.search(
|
||||
query_embedding=embedding,
|
||||
limit=10,
|
||||
filter='doc_id LIKE "ai%" AND chunk_id > 0',
|
||||
)
|
||||
```
|
||||
|
||||
### Table Structure
|
||||
|
||||
The implementation automatically creates a table with the following structure:
|
||||
|
||||
```sql
|
||||
CREATE TABLE IF NOT EXISTS table_name (
|
||||
id VARCHAR(255) PRIMARY KEY,
|
||||
embedding VECTOR(dimensions) NOT NULL,
|
||||
doc_id VARCHAR(255) NOT NULL,
|
||||
chunk_id INT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
total_chunks INT NOT NULL,
|
||||
INDEX idx_doc_id (doc_id),
|
||||
INDEX idx_chunk_id (chunk_id),
|
||||
VECTOR INDEX (embedding) M=16 DISTANCE=cosine -- or DISTANCE=euclidean
|
||||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
|
||||
```
|
||||
|
||||
**Note**: The vector index is created directly within the `CREATE TABLE` statement, not as a separate SQL command. The `M` parameter controls the HNSW algorithm's graph connectivity (default: 16).
|
||||
|
||||
### Performance Considerations
|
||||
|
||||
- **VECTOR Data Type**: Uses MySQL's native VECTOR type for efficient storage
|
||||
- **Vector Index**: Automatically creates a vector index with the specified distance metric for fast similarity search
|
||||
- **Database-Level Distance Calculation**: Vector distance calculations are performed at the database level using MySQL's native vector functions (VEC_DISTANCE_COSINE, VEC_DISTANCE_EUCLIDEAN), with sorting done via SQL ORDER BY
|
||||
- **Native Vector Support**: MySQL 8.0+ has built-in vector functions that are highly optimized for vector operations
|
||||
- **Supported Distance Metrics**: Only COSINE and EUCLIDEAN are supported
|
||||
- **Small to Medium Datasets**: AlibabaCloudMySQLStore performs well for datasets up to 100K vectors
|
||||
- **Large Datasets**: For datasets with millions of vectors, consider using dedicated vector databases (MilvusLite, Qdrant) with specialized indexing (HNSW, IVF, etc.)
|
||||
- **RDS Performance**: Leverage AlibabaCloud RDS features like read replicas, backup, and monitoring
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
### Direct Database Access
|
||||
|
||||
```python
|
||||
# Get the MySQL connection for advanced operations
|
||||
conn = store.get_client()
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Execute custom SQL queries
|
||||
cursor.execute("SELECT COUNT(*) FROM test_vectors")
|
||||
count = cursor.fetchone()
|
||||
print(f"Total vectors: {count[0]}")
|
||||
```
|
||||
|
||||
### Using MySQL Native Vector Functions
|
||||
|
||||
MySQL's native vector functions can be used directly in SQL queries:
|
||||
|
||||
```python
|
||||
conn = store.get_client()
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Use MySQL native vector functions directly
|
||||
query_vector = "[0.1,0.2,0.3,0.4]"
|
||||
cursor.execute("""
|
||||
SELECT
|
||||
doc_id,
|
||||
VEC_DISTANCE_COSINE(vector, VEC_FROMTEXT(%s)) as distance
|
||||
FROM test_vectors
|
||||
ORDER BY distance ASC
|
||||
LIMIT 10
|
||||
""", (query_vector,))
|
||||
|
||||
results = cursor.fetchall()
|
||||
|
||||
# Available MySQL vector functions in AlibabaCloud:
|
||||
# - VEC_FROMTEXT(text) - Convert text "[1,2,3]" to vector
|
||||
# - VEC_DISTANCE_COSINE(v1, v2) - Cosine distance
|
||||
# - VEC_DISTANCE_EUCLIDEAN(v1, v2) - Euclidean distance
|
||||
```
|
||||
|
||||
### SSL Connection
|
||||
|
||||
For secure connections to AlibabaCloud RDS:
|
||||
|
||||
```python
|
||||
store = AlibabaCloudMySQLStore(
|
||||
host="rm-xxxxx.mysql.rds.aliyuncs.com",
|
||||
port=3306,
|
||||
user="your_username",
|
||||
password="your_password",
|
||||
database="agentscope_test",
|
||||
table_name="vectors",
|
||||
dimensions=768,
|
||||
distance="COSINE",
|
||||
connection_kwargs={
|
||||
"ssl_ca": "/path/to/ca.pem",
|
||||
"ssl_verify_cert": True,
|
||||
"ssl_verify_identity": True,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
### Batch Operations
|
||||
|
||||
```python
|
||||
# Add large batches of documents
|
||||
batch_size = 1000
|
||||
for i in range(0, len(all_documents), batch_size):
|
||||
batch = all_documents[i:i + batch_size]
|
||||
await store.add(batch)
|
||||
```
|
||||
|
||||
### Connection Pooling
|
||||
|
||||
```python
|
||||
store = AlibabaCloudMySQLStore(
|
||||
host="rm-xxxxx.mysql.rds.aliyuncs.com",
|
||||
port=3306,
|
||||
user="your_username",
|
||||
password="your_password",
|
||||
database="agentscope_test",
|
||||
table_name="vectors",
|
||||
dimensions=768,
|
||||
distance="COSINE",
|
||||
connection_kwargs={
|
||||
"pool_name": "mypool",
|
||||
"pool_size": 10,
|
||||
"pool_reset_session": True,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### MySQL Version Check
|
||||
|
||||
Ensure your RDS MySQL version supports vector functions:
|
||||
|
||||
```sql
|
||||
SELECT VERSION();
|
||||
-- Should be MySQL 8.0 or higher
|
||||
|
||||
-- Check if vector capability is enabled (critical check)
|
||||
SHOW VARIABLES LIKE 'vidx_disabled';
|
||||
-- Expected result: vidx_disabled | OFF (vector capability enabled)
|
||||
|
||||
-- Test vector functions
|
||||
SELECT VEC_FROMTEXT('[1,2,3]');
|
||||
```
|
||||
|
||||
### Connection Errors
|
||||
|
||||
If you get connection errors:
|
||||
|
||||
1. **Check Whitelist**: Ensure your IP is in the RDS whitelist
|
||||
2. **Security Group**: Verify security group rules allow port 3306
|
||||
3. **Network Type**: Ensure you're using the correct endpoint (public/private)
|
||||
4. **Credentials**: Double-check username and password
|
||||
|
||||
```bash
|
||||
# Test connection from command line
|
||||
mysql -h rm-xxxxx.mysql.rds.aliyuncs.com -P 3306 -u your_username -p
|
||||
```
|
||||
|
||||
### Vector Function Errors
|
||||
|
||||
If you get errors about VEC_DISTANCE_COSINE or VECTOR type not being recognized:
|
||||
|
||||
1. **Check if vector capability is enabled**:
|
||||
|
||||
```sql
|
||||
-- Check vidx_disabled parameter (must be OFF)
|
||||
SHOW VARIABLES LIKE 'vidx_disabled';
|
||||
-- Expected result: vidx_disabled | OFF
|
||||
-- If ON, vector capability is disabled, contact AlibabaCloud support
|
||||
```
|
||||
|
||||
2. Verify MySQL version is 8.0 or higher
|
||||
|
||||
```sql
|
||||
SELECT VERSION();
|
||||
```
|
||||
|
||||
3. Test vector functions availability:
|
||||
|
||||
```sql
|
||||
-- Check if vector functions are available
|
||||
SELECT VEC_FROMTEXT('[1,2,3]');
|
||||
|
||||
-- Check if VECTOR type is supported
|
||||
CREATE TABLE test_vector (v VECTOR(3));
|
||||
DROP TABLE test_vector;
|
||||
|
||||
-- List vector indexes
|
||||
SHOW INDEX FROM your_table_name WHERE Index_type = 'VECTOR';
|
||||
```
|
||||
|
||||
If `vidx_disabled` is ON, contact AlibabaCloud support to enable the vector search plugin for your RDS instance.
|
||||
|
||||
### Performance Optimization
|
||||
|
||||
For large datasets on AlibabaCloud RDS:
|
||||
|
||||
1. **Upgrade Instance**: Consider higher specifications (CPU, Memory)
|
||||
2. **Read Replicas**: Use read replicas for read-heavy workloads
|
||||
3. **Indexes**: Add indexes on frequently filtered columns
|
||||
4. **Connection Pool**: Use connection pooling for concurrent operations
|
||||
5. **Monitor**: Use AlibabaCloud CloudMonitor for performance insights
|
||||
|
||||
### Timeout Errors
|
||||
|
||||
If you experience timeout errors:
|
||||
|
||||
```python
|
||||
store = AlibabaCloudMySQLStore(
|
||||
host="rm-xxxxx.mysql.rds.aliyuncs.com",
|
||||
port=3306,
|
||||
user="your_username",
|
||||
password="your_password",
|
||||
database="agentscope_test",
|
||||
table_name="vectors",
|
||||
dimensions=768,
|
||||
distance="COSINE",
|
||||
connection_kwargs={
|
||||
"connect_timeout": 30,
|
||||
"read_timeout": 60,
|
||||
"write_timeout": 60,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
## AlibabaCloud RDS Best Practices
|
||||
|
||||
1. **Backup**: Enable automatic backups in RDS console
|
||||
2. **Monitoring**: Set up alerts for CPU, memory, and connection usage
|
||||
3. **Security**: Use private network connections when possible
|
||||
4. **Scaling**: Consider read-only instances for read-heavy workloads
|
||||
5. **Cost Optimization**: Use reserved instances for long-term usage
|
||||
|
||||
## Related Resources
|
||||
|
||||
- [AlibabaCloud RDS Documentation](https://www.alibabacloud.com/help/en/apsaradb-for-rds)
|
||||
- [AlibabaCloud MySQL Vector Functions](https://www.alibabacloud.com/help/en/rds/apsaradb-rds-for-mysql/vector-storage-1)
|
||||
- [AgentScope RAG Tutorial](https://doc.agentscope.io/tutorial/task_rag.html)
|
||||
- [MySQL Connector Python](https://dev.mysql.com/doc/connector-python/en/)
|
||||
|
||||
## Example Use Cases
|
||||
|
||||
### RAG System with AlibabaCloud
|
||||
|
||||
```python
|
||||
from agentscope.rag import AlibabaCloudMySQLStore, KnowledgeBase
|
||||
|
||||
# Initialize vector store with AlibabaCloud RDS
|
||||
store = AlibabaCloudMySQLStore(
|
||||
host="rm-xxxxx.mysql.rds.aliyuncs.com",
|
||||
port=3306,
|
||||
user="your_username",
|
||||
password="your_password",
|
||||
database="rag_system",
|
||||
table_name="knowledge_vectors",
|
||||
dimensions=768,
|
||||
distance="COSINE",
|
||||
)
|
||||
|
||||
# Create knowledge base
|
||||
kb = KnowledgeBase(store=store)
|
||||
|
||||
# Add documents
|
||||
await kb.add_documents(documents)
|
||||
|
||||
# Search
|
||||
results = await kb.search("What is AI?", top_k=5)
|
||||
```
|
||||
|
||||
## Support
|
||||
|
||||
For issues related to:
|
||||
- **AlibabaCloudMySQLStore**: Open an issue on AgentScope GitHub
|
||||
- **RDS MySQL**: Contact AlibabaCloud Support
|
||||
- **Vector Functions**: Check MySQL documentation or AlibabaCloud support
|
||||
|
||||
## License
|
||||
|
||||
This example is part of the AgentScope project and follows the same license.
|
||||
|
||||
@@ -0,0 +1,282 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Example of using AlibabaCloudMySQLStore in AgentScope RAG system."""
|
||||
import asyncio
|
||||
from agentscope.rag import (
|
||||
AlibabaCloudMySQLStore,
|
||||
Document,
|
||||
DocMetadata,
|
||||
)
|
||||
from agentscope.message import TextBlock
|
||||
|
||||
|
||||
async def example_basic_operations() -> None:
|
||||
"""The example of basic CRUD operations with AlibabaCloudMySQLStore."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 1: Basic CRUD Operations")
|
||||
print("=" * 60)
|
||||
|
||||
# Initialize AlibabaCloudMySQLStore
|
||||
# Replace with your AlibabaCloud MySQL connection details
|
||||
store = AlibabaCloudMySQLStore(
|
||||
host="rm-xxxxx.mysql.rds.aliyuncs.com", # Your RDS endpoint
|
||||
port=3306,
|
||||
user="your_username",
|
||||
password="your_password",
|
||||
database="agentscope_test",
|
||||
table_name="test_vectors",
|
||||
dimensions=4, # Small dimension for testing
|
||||
distance="COSINE",
|
||||
)
|
||||
|
||||
print("✓ AlibabaCloudMySQLStore initialized")
|
||||
|
||||
# Create test documents with embeddings
|
||||
test_docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
text="Artificial Intelligence is the future",
|
||||
),
|
||||
doc_id="doc_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Machine Learning is a subset of AI"),
|
||||
doc_id="doc_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Deep Learning uses neural networks"),
|
||||
doc_id="doc_3",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
]
|
||||
|
||||
# Test add operation
|
||||
await store.add(test_docs)
|
||||
print(f"✓ Added {len(test_docs)} documents to the store")
|
||||
|
||||
# Test search operation
|
||||
query_embedding = [0.15, 0.25, 0.35, 0.45]
|
||||
results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=2,
|
||||
)
|
||||
|
||||
print(f"\n✓ Search completed, found {len(results)} results:")
|
||||
for i, result in enumerate(results, 1):
|
||||
print(f" {i}. Score: {result.score:.4f}")
|
||||
print(f" Content: {result.metadata.content}")
|
||||
print(f" Doc ID: {result.metadata.doc_id}")
|
||||
|
||||
# Test search with score threshold
|
||||
results_filtered = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=5,
|
||||
score_threshold=0.9,
|
||||
)
|
||||
print(f"\n✓ Search with threshold (>0.9): {len(results_filtered)} results")
|
||||
|
||||
# Test delete operation
|
||||
await store.delete(filter='doc_id = "doc_2"')
|
||||
print("\n✓ Deleted document with doc_id='doc_2'")
|
||||
|
||||
# Verify deletion
|
||||
results_after_delete = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=5,
|
||||
)
|
||||
print(f"✓ After deletion: {len(results_after_delete)} documents remain")
|
||||
|
||||
# Get client for advanced operations
|
||||
client = store.get_client()
|
||||
print(f"\n✓ Got MySQL connection: {type(client).__name__}")
|
||||
|
||||
# Close connection
|
||||
store.close()
|
||||
print("✓ Connection closed")
|
||||
|
||||
|
||||
async def example_filter_search() -> None:
|
||||
"""The example of search with metadata filtering."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 2: Search with Metadata Filtering")
|
||||
print("=" * 60)
|
||||
|
||||
store = AlibabaCloudMySQLStore(
|
||||
host="rm-xxxxx.mysql.rds.aliyuncs.com",
|
||||
port=3306,
|
||||
user="your_username",
|
||||
password="your_password",
|
||||
database="agentscope_test",
|
||||
table_name="filter_vectors",
|
||||
dimensions=4,
|
||||
distance="COSINE",
|
||||
)
|
||||
|
||||
# Create documents with different categories
|
||||
docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Python is a programming language"),
|
||||
doc_id="prog_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
text="Java is used for enterprise applications",
|
||||
),
|
||||
doc_id="prog_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Neural networks are used in AI"),
|
||||
doc_id="ai_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Deep learning requires GPUs"),
|
||||
doc_id="ai_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.4, 0.5, 0.6, 0.7],
|
||||
),
|
||||
]
|
||||
|
||||
await store.add(docs)
|
||||
print(f"✓ Added {len(docs)} documents with different doc_id prefixes")
|
||||
|
||||
# Search without filter
|
||||
query_embedding = [0.25, 0.35, 0.45, 0.55]
|
||||
all_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
)
|
||||
print(f"\n✓ Search without filter: {len(all_results)} results")
|
||||
for i, result in enumerate(all_results, 1):
|
||||
doc_id = result.metadata.doc_id
|
||||
score = result.score
|
||||
print(f" {i}. Doc ID: {doc_id}, Score: {score:.4f}")
|
||||
|
||||
# Search with filter for programming docs
|
||||
prog_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
filter='doc_id LIKE "prog%"',
|
||||
)
|
||||
filter_msg = 'doc_id LIKE "prog%"'
|
||||
print(f"\n✓ Search with filter ({filter_msg}): {len(prog_results)}")
|
||||
for i, result in enumerate(prog_results, 1):
|
||||
doc_id = result.metadata.doc_id
|
||||
score = result.score
|
||||
print(f" {i}. Doc ID: {doc_id}, Score: {score:.4f}")
|
||||
|
||||
# Search with filter for AI docs
|
||||
ai_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
filter='doc_id LIKE "ai%"',
|
||||
)
|
||||
filter_msg = 'doc_id LIKE "ai%"'
|
||||
print(f"\n✓ Search with filter ({filter_msg}): {len(ai_results)}")
|
||||
for i, result in enumerate(ai_results, 1):
|
||||
doc_id = result.metadata.doc_id
|
||||
score = result.score
|
||||
print(f" {i}. Doc ID: {doc_id}, Score: {score:.4f}")
|
||||
|
||||
store.close()
|
||||
|
||||
|
||||
async def example_distance_metrics() -> None:
|
||||
"""The example of different distance metrics."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 3: Different Distance Metrics")
|
||||
print("=" * 60)
|
||||
|
||||
# Test with different metrics
|
||||
# Note: AlibabaCloud MySQL only supports COSINE and EUCLIDEAN
|
||||
metrics = ["COSINE", "EUCLIDEAN"]
|
||||
|
||||
for metric in metrics:
|
||||
print(f"\n--- Testing {metric} metric ---")
|
||||
store = AlibabaCloudMySQLStore(
|
||||
host="rm-xxxxx.mysql.rds.aliyuncs.com",
|
||||
port=3306,
|
||||
user="your_username",
|
||||
password="your_password",
|
||||
database="agentscope_test",
|
||||
table_name=f"{metric.lower()}_vectors",
|
||||
dimensions=4,
|
||||
distance=metric,
|
||||
)
|
||||
|
||||
docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text=f"Test doc for {metric}"),
|
||||
doc_id=f"doc_{metric}_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
]
|
||||
|
||||
await store.add(docs)
|
||||
results = await store.search(
|
||||
query_embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
limit=1,
|
||||
)
|
||||
|
||||
print(f"✓ {metric} metric: Score = {results[0].score:.4f}")
|
||||
store.close()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Run all examples."""
|
||||
print("\n" + "=" * 60)
|
||||
print("AlibabaCloud MySQL Vector Store Test Suite")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
await example_basic_operations()
|
||||
await example_filter_search()
|
||||
await example_distance_metrics()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✓ All tests completed successfully!")
|
||||
print("=" * 60)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n✗ Test failed with error: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
128
examples/functionality/vector_store/milvus_lite/README.md
Normal file
128
examples/functionality/vector_store/milvus_lite/README.md
Normal file
@@ -0,0 +1,128 @@
|
||||
# MilvusLite Vector Store
|
||||
|
||||
This example demonstrates how to use **MilvusLiteStore** for vector storage and semantic search in AgentScope.
|
||||
It includes four test scenarios covering CRUD operations, metadata filtering, document chunking, and distance metrics.
|
||||
|
||||
### Quick Start
|
||||
|
||||
Install agentscope first, and then the MilvusLite dependency:
|
||||
|
||||
```bash
|
||||
# In MacOS/Linux
|
||||
pip install pymilvus\[milvus_lite\]
|
||||
|
||||
# In Windows
|
||||
pip install pymilvus[milvus_lite]
|
||||
```
|
||||
|
||||
Run the example script, which showcases adding, searching with/without filters in MilvusLite vector store:
|
||||
|
||||
```bash
|
||||
python milvuslite_store.py
|
||||
```
|
||||
|
||||
> **Note:** The script creates `.db` files in the current directory. You can delete them after testing.
|
||||
|
||||
## Usage
|
||||
|
||||
### Initialize Store
|
||||
```python
|
||||
from agentscope.rag import MilvusLiteStore
|
||||
|
||||
store = MilvusLiteStore(
|
||||
uri="./milvus_test.db",
|
||||
collection_name="test_collection",
|
||||
dimensions=768, # Match your embedding model
|
||||
distance="COSINE", # COSINE, L2, or IP
|
||||
)
|
||||
```
|
||||
|
||||
### Add Documents
|
||||
|
||||
```python
|
||||
from agentscope.rag import Document, DocMetadata
|
||||
from agentscope.message import TextBlock
|
||||
|
||||
doc = Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(type="text", text="Your document text"),
|
||||
doc_id="doc_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, ...], # Your embedding vector
|
||||
)
|
||||
|
||||
await store.add([doc])
|
||||
```
|
||||
|
||||
### Search
|
||||
|
||||
```python
|
||||
results = await store.search(
|
||||
query_embedding=[0.15, 0.25, ...],
|
||||
limit=5,
|
||||
score_threshold=0.9, # Optional
|
||||
filter='doc_id like "prefix%"', # Optional
|
||||
)
|
||||
```
|
||||
|
||||
### Delete
|
||||
|
||||
```python
|
||||
await store.delete(filter_expr='doc_id == "doc_1"')
|
||||
```
|
||||
|
||||
## Distance Metrics
|
||||
|
||||
| Metric | Description | Best For |
|
||||
|--------|-------------|----------|
|
||||
| **COSINE** | Cosine similarity | Text embeddings (recommended) |
|
||||
| **L2** | Euclidean distance | Spatial data |
|
||||
| **IP** | Inner Product | Recommendation systems |
|
||||
|
||||
## Filter Expressions
|
||||
|
||||
```python
|
||||
# Exact match
|
||||
filter='doc_id == "doc_1"'
|
||||
|
||||
# Pattern matching
|
||||
filter='doc_id like "prefix%"'
|
||||
|
||||
# Numeric and logical operators
|
||||
filter='chunk_id >= 0 and total_chunks > 1'
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
### Access Underlying Client
|
||||
```python
|
||||
client = store.get_client()
|
||||
stats = client.get_collection_stats(collection_name="test_collection")
|
||||
```
|
||||
|
||||
### Document Metadata
|
||||
- `content`: Text content (TextBlock)
|
||||
- `doc_id`: Unique document identifier
|
||||
- `chunk_id`: Chunk position (0-indexed)
|
||||
- `total_chunks`: Total chunks in document
|
||||
|
||||
## FAQ
|
||||
|
||||
**What embedding dimension should I use?**
|
||||
Match your embedding model's output dimension (e.g., 768 for BERT, 1536 for OpenAI ada-002).
|
||||
|
||||
**Can I change the distance metric after creation?**
|
||||
No, create a new collection with the desired metric.
|
||||
|
||||
**How do I delete the database?**
|
||||
Delete the `.db` file specified in the `uri` parameter.
|
||||
|
||||
**Is this suitable for production?**
|
||||
MilvusLite works well for development and small-scale applications. For production at scale, consider Milvus standalone or cluster mode.
|
||||
|
||||
## References
|
||||
|
||||
- [Milvus Documentation](https://milvus.io/docs)
|
||||
- [AgentScope RAG Tutorial](https://doc.agentscope.io/tutorial/task_rag.html)
|
||||
327
examples/functionality/vector_store/milvus_lite/main.py
Normal file
327
examples/functionality/vector_store/milvus_lite/main.py
Normal file
@@ -0,0 +1,327 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Example of using MilvusLiteStore in AgentScope RAG system."""
|
||||
import asyncio
|
||||
from agentscope.rag import (
|
||||
MilvusLiteStore,
|
||||
Document,
|
||||
DocMetadata,
|
||||
)
|
||||
from agentscope.message import TextBlock
|
||||
|
||||
|
||||
async def example_basic_operations() -> None:
|
||||
"""The example of basic CRUD operations with MilvusLiteStore."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 1: Basic CRUD Operations")
|
||||
print("=" * 60)
|
||||
|
||||
# Initialize MilvusLiteStore with a local file
|
||||
store = MilvusLiteStore(
|
||||
uri="./milvus_test.db",
|
||||
collection_name="test_collection",
|
||||
dimensions=4, # Small dimension for testing
|
||||
distance="COSINE",
|
||||
)
|
||||
|
||||
print("✓ MilvusLiteStore initialized")
|
||||
|
||||
# Create test documents with embeddings
|
||||
test_docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
text="Artificial Intelligence is the future",
|
||||
),
|
||||
doc_id="doc_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Machine Learning is a subset of AI"),
|
||||
doc_id="doc_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Deep Learning uses neural networks"),
|
||||
doc_id="doc_3",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
]
|
||||
|
||||
# Test add operation
|
||||
await store.add(test_docs)
|
||||
print(f"✓ Added {len(test_docs)} documents to the store")
|
||||
|
||||
# Test search operation
|
||||
query_embedding = [0.15, 0.25, 0.35, 0.45]
|
||||
results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=2,
|
||||
)
|
||||
|
||||
print(f"\n✓ Search completed, found {len(results)} results:")
|
||||
for i, result in enumerate(results, 1):
|
||||
print(f" {i}. Score: {result.score:.4f}")
|
||||
print(f" Content: {result.metadata.content}")
|
||||
print(f" Doc ID: {result.metadata.doc_id}")
|
||||
|
||||
# Test search with score threshold
|
||||
results_filtered = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=5,
|
||||
score_threshold=0.9,
|
||||
)
|
||||
print(f"\n✓ Search with threshold (>0.9): {len(results_filtered)} results")
|
||||
|
||||
# Test delete operation
|
||||
# Note: We need to use filter expression to delete by doc_id
|
||||
await store.delete(filter='doc_id == "doc_2"')
|
||||
print("\n✓ Deleted document with doc_id='doc_2'")
|
||||
|
||||
# Verify deletion
|
||||
results_after_delete = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=5,
|
||||
)
|
||||
print(f"✓ After deletion: {len(results_after_delete)} documents remain")
|
||||
|
||||
# Get client for advanced operations
|
||||
client = store.get_client()
|
||||
print(f"\n✓ Got MilvusClient: {type(client).__name__}")
|
||||
|
||||
|
||||
async def example_filter_search() -> None:
|
||||
"""The example of search with metadata filtering."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 2: Search with Metadata Filtering")
|
||||
print("=" * 60)
|
||||
|
||||
store = MilvusLiteStore(
|
||||
uri="./milvus_filter_test.db",
|
||||
collection_name="filter_collection",
|
||||
dimensions=4,
|
||||
distance="COSINE",
|
||||
)
|
||||
|
||||
# Create documents with different categories
|
||||
docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Python is a programming language"),
|
||||
doc_id="prog_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
text="Java is used for enterprise applications",
|
||||
),
|
||||
doc_id="prog_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Neural networks are used in AI"),
|
||||
doc_id="ai_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Deep learning requires GPUs"),
|
||||
doc_id="ai_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.4, 0.5, 0.6, 0.7],
|
||||
),
|
||||
]
|
||||
|
||||
await store.add(docs)
|
||||
print(f"✓ Added {len(docs)} documents with different doc_id prefixes")
|
||||
|
||||
# Search without filter
|
||||
query_embedding = [0.25, 0.35, 0.45, 0.55]
|
||||
all_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
)
|
||||
print(f"\n✓ Search without filter: {len(all_results)} results")
|
||||
for i, result in enumerate(all_results, 1):
|
||||
doc_id = result.metadata.doc_id
|
||||
score = result.score
|
||||
print(f" {i}. Doc ID: {doc_id}, Score: {score:.4f}")
|
||||
|
||||
# Search with filter for programming docs
|
||||
prog_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
filter='doc_id like "prog%"',
|
||||
)
|
||||
filter_msg = "doc_id like 'prog%'"
|
||||
print(f"\n✓ Search with filter ({filter_msg}): {len(prog_results)}")
|
||||
for i, result in enumerate(prog_results, 1):
|
||||
doc_id = result.metadata.doc_id
|
||||
score = result.score
|
||||
print(f" {i}. Doc ID: {doc_id}, Score: {score:.4f}")
|
||||
|
||||
# Search with filter for AI docs
|
||||
ai_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
filter='doc_id like "ai%"',
|
||||
)
|
||||
filter_msg = "doc_id like 'ai%'"
|
||||
print(f"\n✓ Search with filter ({filter_msg}): {len(ai_results)}")
|
||||
for i, result in enumerate(ai_results, 1):
|
||||
doc_id = result.metadata.doc_id
|
||||
score = result.score
|
||||
print(f" {i}. Doc ID: {doc_id}, Score: {score:.4f}")
|
||||
|
||||
|
||||
async def example_multiple_chunks() -> None:
|
||||
"""The example of documents with multiple chunks."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 3: Documents with Multiple Chunks")
|
||||
print("=" * 60)
|
||||
|
||||
store = MilvusLiteStore(
|
||||
uri="./milvus_chunks_test.db",
|
||||
collection_name="chunks_collection",
|
||||
dimensions=4,
|
||||
distance="COSINE",
|
||||
)
|
||||
|
||||
# Create a document split into multiple chunks
|
||||
chunks = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Chapter 1: Introduction to AI"),
|
||||
doc_id="book_1",
|
||||
chunk_id=0,
|
||||
total_chunks=3,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Chapter 2: Machine Learning Basics"),
|
||||
doc_id="book_1",
|
||||
chunk_id=1,
|
||||
total_chunks=3,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Chapter 3: Deep Learning Advanced"),
|
||||
doc_id="book_1",
|
||||
chunk_id=2,
|
||||
total_chunks=3,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
]
|
||||
|
||||
await store.add(chunks)
|
||||
print(f"✓ Added document with {len(chunks)} chunks")
|
||||
|
||||
# Search and verify chunk information
|
||||
query_embedding = [0.2, 0.3, 0.4, 0.5]
|
||||
results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=3,
|
||||
)
|
||||
|
||||
print("\n✓ Search results for multi-chunk document:")
|
||||
for i, result in enumerate(results, 1):
|
||||
chunk_info = (
|
||||
f"{result.metadata.chunk_id}/{result.metadata.total_chunks}"
|
||||
)
|
||||
print(f" {i}. Chunk {chunk_info}")
|
||||
print(f" Content: {result.metadata.content}")
|
||||
print(f" Score: {result.score:.4f}")
|
||||
|
||||
|
||||
async def example_distance_metrics() -> None:
|
||||
"""The example of different distance metrics."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 4: Different Distance Metrics")
|
||||
print("=" * 60)
|
||||
|
||||
# Test with different metrics
|
||||
metrics = ["COSINE", "L2", "IP"]
|
||||
|
||||
for metric in metrics:
|
||||
print(f"\n--- Testing {metric} metric ---")
|
||||
store = MilvusLiteStore(
|
||||
uri=f"./milvus_{metric.lower()}_test.db",
|
||||
collection_name=f"{metric.lower()}_collection",
|
||||
dimensions=4,
|
||||
distance=metric,
|
||||
)
|
||||
|
||||
docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text=f"Test doc for {metric}"),
|
||||
doc_id=f"doc_{metric}_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
]
|
||||
|
||||
await store.add(docs)
|
||||
results = await store.search(
|
||||
query_embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
limit=1,
|
||||
)
|
||||
|
||||
print(f"✓ {metric} metric: Score = {results[0].score:.4f}")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Run all example."""
|
||||
print("\n" + "=" * 60)
|
||||
print("MilvusLiteStore Comprehensive Test Suite")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
await example_basic_operations()
|
||||
await example_filter_search()
|
||||
await example_multiple_chunks()
|
||||
await example_distance_metrics()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✓ All tests completed successfully!")
|
||||
print("=" * 60)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n✗ Test failed with error: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
211
examples/functionality/vector_store/mongodb/README.md
Normal file
211
examples/functionality/vector_store/mongodb/README.md
Normal file
@@ -0,0 +1,211 @@
|
||||
# MongoDB Vector Store
|
||||
|
||||
This example demonstrates how to use **MongoDBStore** for vector storage and semantic search in AgentScope using MongoDB's Vector Search capabilities.
|
||||
It includes comprehensive test scenarios covering CRUD operations, metadata filtering, document chunking, and distance metrics.
|
||||
|
||||
### Quick Start
|
||||
|
||||
Install agentscope first, and then the MongoDB dependency:
|
||||
|
||||
```bash
|
||||
pip install pymongo
|
||||
```
|
||||
|
||||
**Important:** Before running the example, you need to set the `MONGODB_HOST`
|
||||
environment variable with your MongoDB connection string:
|
||||
|
||||
```bash
|
||||
# For local MongoDB
|
||||
export MONGODB_HOST="mongodb://localhost:27017/?directConnection=true"
|
||||
|
||||
# For MongoDB Atlas (replace with your connection string)
|
||||
# export MONGODB_HOST=${YOUR_MONGODB_HOST}
|
||||
```
|
||||
|
||||
Run the example script, which showcases adding, searching, and deleting in MongoDB vector store:
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
|
||||
> **Note:** The script connects to MongoDB Atlas or local MongoDB instance. Make sure you have a valid MongoDB connection string.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Confirm your MongoDB instance supports Vector Search functionality
|
||||
- Valid MongoDB connection string (local or Atlas)
|
||||
|
||||
## Usage
|
||||
|
||||
### Initialize Store
|
||||
|
||||
```python
|
||||
from agentscope.rag import MongoDBStore
|
||||
|
||||
# For MongoDB Atlas
|
||||
store = MongoDBStore(
|
||||
host="mongodb+srv://username:password@cluster.mongodb.net/",
|
||||
db_name="test_db",
|
||||
collection_name="test_collection",
|
||||
dimensions=768, # Match your embedding model
|
||||
distance="cosine", # cosine, euclidean, or dotProduct
|
||||
)
|
||||
|
||||
# For local MongoDB
|
||||
store = MongoDBStore(
|
||||
host="mongodb://localhost:27017/?directConnection=true",
|
||||
db_name="test_db",
|
||||
collection_name="test_collection",
|
||||
dimensions=768,
|
||||
distance="cosine",
|
||||
)
|
||||
|
||||
# To enable filtering in search, specify filter_fields:
|
||||
store = MongoDBStore(
|
||||
host="mongodb://localhost:27017/?directConnection=true",
|
||||
db_name="test_db",
|
||||
collection_name="test_collection",
|
||||
dimensions=768,
|
||||
distance="cosine",
|
||||
filter_fields=["payload.doc_id", "payload.chunk_id"], # Fields for filtering
|
||||
)
|
||||
|
||||
# No manual initialization needed - everything is automatic!
|
||||
# Database, collection, and vector search index are created automatically
|
||||
# when you first call add() or search()
|
||||
```
|
||||
|
||||
### Add Documents
|
||||
|
||||
```python
|
||||
from agentscope.rag import Document, DocMetadata
|
||||
from agentscope.message import TextBlock
|
||||
|
||||
doc = Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(type="text", text="Your document text"),
|
||||
doc_id="doc_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, ...], # Your embedding vector
|
||||
)
|
||||
|
||||
await store.add([doc])
|
||||
```
|
||||
|
||||
### Search
|
||||
|
||||
```python
|
||||
results = await store.search(
|
||||
query_embedding=[0.15, 0.25, ...],
|
||||
limit=5,
|
||||
score_threshold=0.9, # Optional
|
||||
filter={"payload.doc_id": {"$in": ["doc_1", "doc_2"]}}, # Optional filter
|
||||
)
|
||||
# Note:
|
||||
# - To use filter, the field must be declared in filter_fields when creating store
|
||||
# - MongoDB $vectorSearch filter supports: $gt, $gte, $lt, $lte,
|
||||
# $eq, $ne, $in, $nin, $exists, $not (NOT $regex)
|
||||
```
|
||||
|
||||
### Delete
|
||||
|
||||
```python
|
||||
# Delete by document IDs (no initialization needed)
|
||||
await store.delete(ids=["doc_1", "doc_2"])
|
||||
|
||||
# Delete entire collection (use with caution)
|
||||
await store.delete_collection()
|
||||
|
||||
# Delete entire database (use with caution)
|
||||
await store.delete_database()
|
||||
```
|
||||
|
||||
## Distance Metrics
|
||||
|
||||
| Metric | Description | Best For |
|
||||
|--------|-------------|----------|
|
||||
| **cosine** | Cosine similarity | Text embeddings (recommended) |
|
||||
| **euclidean** | Euclidean distance | Spatial data |
|
||||
| **dotProduct** | Inner Product | Recommendation systems |
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
### Access Underlying Client
|
||||
|
||||
```python
|
||||
client = store.get_client()
|
||||
# Use MongoDB client for advanced operations
|
||||
stats = await client[store.db_name].command("collStats", store.collection_name)
|
||||
```
|
||||
|
||||
### Document Metadata
|
||||
|
||||
- `content`: Text content (TextBlock)
|
||||
- `doc_id`: Unique document identifier
|
||||
- `chunk_id`: Chunk position (0-indexed)
|
||||
- `total_chunks`: Total chunks in document
|
||||
|
||||
### Vector Search Index
|
||||
|
||||
MongoDBStore automatically creates vector search indexes with the following configuration:
|
||||
|
||||
```python
|
||||
{
|
||||
"fields": [
|
||||
{
|
||||
"type": "vector",
|
||||
"path": "vector",
|
||||
"similarity": "cosine", # or euclidean, dotProduct
|
||||
"numDimensions": 768
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
## Connection Examples
|
||||
|
||||
### MongoDB Atlas
|
||||
|
||||
```python
|
||||
store = MongoDBStore(
|
||||
host="<YOUR_MONGO_ATLAS_CONNECTION_STRING>",
|
||||
db_name="production_db",
|
||||
collection_name="documents",
|
||||
dimensions=1536,
|
||||
distance="cosine",
|
||||
)
|
||||
```
|
||||
|
||||
### Local MongoDB
|
||||
|
||||
#### Without Authentication
|
||||
|
||||
```python
|
||||
store = MongoDBStore(
|
||||
host="mongodb://localhost:27017?directConnection=true",
|
||||
db_name="local_db",
|
||||
collection_name="test_collection",
|
||||
dimensions=768,
|
||||
distance="cosine",
|
||||
)
|
||||
```
|
||||
|
||||
#### With Authentication
|
||||
|
||||
```python
|
||||
store = MongoDBStore(
|
||||
host="mongodb://user:pass@localhost:27017/?directConnection=true",
|
||||
db_name="test_db",
|
||||
collection_name="test_collection",
|
||||
dimensions=768,
|
||||
distance="cosine",
|
||||
)
|
||||
```
|
||||
|
||||
## References
|
||||
|
||||
- [MongoDB Vector Search Documentation](https://www.mongodb.com/docs/atlas/atlas-search/vector-search/)
|
||||
- [MongoDB Atlas Documentation](https://www.mongodb.com/docs/atlas/)
|
||||
- [AgentScope RAG Tutorial](https://doc.agentscope.io/tutorial/task_rag.html)
|
||||
351
examples/functionality/vector_store/mongodb/main.py
Normal file
351
examples/functionality/vector_store/mongodb/main.py
Normal file
@@ -0,0 +1,351 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Example of using MongoDBStore in AgentScope RAG system."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.rag import (
|
||||
MongoDBStore,
|
||||
Document,
|
||||
DocMetadata,
|
||||
)
|
||||
from agentscope.message import TextBlock
|
||||
|
||||
|
||||
async def example_basic_operations() -> None:
|
||||
"""The example of basic CRUD operations with MongoDBStore."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 1: Basic CRUD Operations")
|
||||
print("=" * 60)
|
||||
|
||||
# Initialize MongoDBStore with MongoDB connection
|
||||
store = MongoDBStore(
|
||||
host=os.getenv("MONGODB_HOST"),
|
||||
db_name="test_db",
|
||||
collection_name="test_collection",
|
||||
dimensions=4, # Small dimension for testing
|
||||
distance="cosine",
|
||||
)
|
||||
|
||||
print("✓ MongoDBStore initialized")
|
||||
|
||||
# Create test documents with embeddings
|
||||
test_docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
text="Artificial Intelligence is the future",
|
||||
),
|
||||
doc_id="doc_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Machine Learning is a subset of AI"),
|
||||
doc_id="doc_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Deep Learning uses neural networks"),
|
||||
doc_id="doc_3",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
]
|
||||
|
||||
# Test add operation (automatically creates database, collection,
|
||||
# and index)
|
||||
await store.add(test_docs)
|
||||
print(f"✓ Added {len(test_docs)} documents to the store")
|
||||
|
||||
# Test search operation (automatically waits for index to be ready)
|
||||
query_embedding = [0.15, 0.25, 0.35, 0.45]
|
||||
results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=2,
|
||||
)
|
||||
|
||||
print(f"\n✓ Search completed, found {len(results)} results:")
|
||||
for i, result in enumerate(results, 1):
|
||||
print(f" {i}. Score: {result.score:.4f}")
|
||||
print(f" Content: {result.metadata.content}")
|
||||
print(f" Doc ID: {result.metadata.doc_id}")
|
||||
|
||||
# Test search with score threshold (also waits for index if needed)
|
||||
results_filtered = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=5,
|
||||
score_threshold=0.3,
|
||||
)
|
||||
print(f"\n✓ Search with threshold (>0.3): {len(results_filtered)} results")
|
||||
|
||||
# Test delete operation (no initialization needed)
|
||||
# Note: MongoDBStore uses ids parameter for deletion
|
||||
await store.delete(ids=["doc_2", "doc_3", "doc_1"])
|
||||
print("\n✓ Deleted documents with specified doc_ids")
|
||||
|
||||
# Verify deletion (search will wait for index if needed)
|
||||
results_after_delete = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=5,
|
||||
)
|
||||
print(f"✓ After deletion: {len(results_after_delete)} documents remain")
|
||||
|
||||
# Get client for advanced operations
|
||||
client = store.get_client()
|
||||
print(f"\n✓ Got MongoDB Client: {type(client).__name__}")
|
||||
|
||||
await store.close()
|
||||
|
||||
|
||||
async def example_filter_search() -> None:
|
||||
"""The example of search with metadata filtering."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 2: Search with Metadata Filtering")
|
||||
print("=" * 60)
|
||||
|
||||
# To use filter in search, specify filter_fields when creating the store.
|
||||
# These fields will be indexed for filtering in $vectorSearch.
|
||||
store = MongoDBStore(
|
||||
host=os.getenv("MONGODB_HOST"),
|
||||
db_name="filter_test_db",
|
||||
collection_name="filter_collection",
|
||||
dimensions=4,
|
||||
distance="cosine",
|
||||
filter_fields=["payload.doc_id"], # Enable filtering on doc_id
|
||||
)
|
||||
|
||||
# Create documents with different categories
|
||||
docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Python is a programming language"),
|
||||
doc_id="prog_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
text="Java is used for enterprise applications",
|
||||
),
|
||||
doc_id="prog_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Neural networks are used in AI"),
|
||||
doc_id="ai_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Deep learning requires GPUs"),
|
||||
doc_id="ai_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.4, 0.5, 0.6, 0.7],
|
||||
),
|
||||
]
|
||||
|
||||
# Add documents (automatically creates database, collection, and index)
|
||||
await store.add(docs)
|
||||
print(f"✓ Added {len(docs)} documents with different doc_id prefixes")
|
||||
|
||||
# Search without filter (automatically waits for index if needed)
|
||||
query_embedding = [0.25, 0.35, 0.45, 0.55]
|
||||
all_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
)
|
||||
print(f"\n✓ Search without filter: {len(all_results)} results")
|
||||
for i, result in enumerate(all_results, 1):
|
||||
doc_id = result.metadata.doc_id
|
||||
score = result.score
|
||||
print(f" {i}. Doc ID: {doc_id}, Score: {score:.4f}")
|
||||
|
||||
# Search with filter for programming docs
|
||||
# Note: doc_id is stored in payload.doc_id in MongoDB documents
|
||||
# MongoDB $vectorSearch filter supports: $gt, $gte, $lt, $lte, $eq, $ne,
|
||||
# $in, $nin, $exists, $not (NOT $regex)
|
||||
prog_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
filter={"payload.doc_id": {"$in": ["prog_1", "prog_2"]}},
|
||||
)
|
||||
print(f"\n✓ Search with filter (prog docs): {len(prog_results)} results")
|
||||
for i, result in enumerate(prog_results, 1):
|
||||
doc_id = result.metadata.doc_id
|
||||
score = result.score
|
||||
print(f" {i}. Doc ID: {doc_id}, Score: {score:.4f}")
|
||||
|
||||
# Search with filter for AI docs
|
||||
ai_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
filter={"payload.doc_id": {"$in": ["ai_1", "ai_2"]}},
|
||||
)
|
||||
print(f"\n✓ Search with filter (ai docs): {len(ai_results)} results")
|
||||
for i, result in enumerate(ai_results, 1):
|
||||
doc_id = result.metadata.doc_id
|
||||
score = result.score
|
||||
print(f" {i}. Doc ID: {doc_id}, Score: {score:.4f}")
|
||||
|
||||
await store.close()
|
||||
|
||||
|
||||
async def example_multiple_chunks() -> None:
|
||||
"""The example of documents with multiple chunks."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 3: Documents with Multiple Chunks")
|
||||
print("=" * 60)
|
||||
|
||||
store = MongoDBStore(
|
||||
host=os.getenv("MONGODB_HOST"),
|
||||
db_name="chunks_test_db",
|
||||
collection_name="chunks_collection",
|
||||
dimensions=4,
|
||||
distance="cosine",
|
||||
)
|
||||
|
||||
# Create a document split into multiple chunks
|
||||
chunks = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Chapter 1: Introduction to AI"),
|
||||
doc_id="book_1",
|
||||
chunk_id=0,
|
||||
total_chunks=3,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Chapter 2: Machine Learning Basics"),
|
||||
doc_id="book_1",
|
||||
chunk_id=1,
|
||||
total_chunks=3,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text="Chapter 3: Deep Learning Advanced"),
|
||||
doc_id="book_1",
|
||||
chunk_id=2,
|
||||
total_chunks=3,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
]
|
||||
|
||||
# Add chunks (automatically creates database, collection, and index)
|
||||
await store.add(chunks)
|
||||
print(f"✓ Added document with {len(chunks)} chunks")
|
||||
|
||||
# Search and verify chunk information (automatically waits for index if
|
||||
# needed)
|
||||
query_embedding = [0.2, 0.3, 0.4, 0.5]
|
||||
results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=3,
|
||||
)
|
||||
|
||||
print("\n✓ Search results for multi-chunk document:")
|
||||
for i, result in enumerate(results, 1):
|
||||
chunk_info = (
|
||||
f"{result.metadata.chunk_id}/{result.metadata.total_chunks}"
|
||||
)
|
||||
print(f" {i}. Chunk {chunk_info}")
|
||||
print(f" Content: {result.metadata.content}")
|
||||
print(f" Score: {result.score:.4f}")
|
||||
|
||||
await store.close()
|
||||
|
||||
|
||||
async def example_distance_metrics() -> None:
|
||||
"""The example of different distance metrics."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 4: Different Distance Metrics")
|
||||
print("=" * 60)
|
||||
|
||||
# Test with different metrics
|
||||
metrics = ["cosine", "euclidean", "dotProduct"]
|
||||
|
||||
for metric in metrics:
|
||||
print(f"\n--- Testing {metric} metric ---")
|
||||
store = MongoDBStore(
|
||||
host=os.getenv("MONGODB_HOST"),
|
||||
db_name=f"{metric}_test_db",
|
||||
collection_name=f"{metric}_collection",
|
||||
dimensions=4,
|
||||
distance=metric,
|
||||
)
|
||||
|
||||
docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(text=f"Test doc for {metric}"),
|
||||
doc_id=f"doc_{metric}_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
]
|
||||
|
||||
# Add and search (automatically creates database/collection/index
|
||||
# and waits for index)
|
||||
await store.add(docs)
|
||||
results = await store.search(
|
||||
query_embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
limit=1,
|
||||
)
|
||||
|
||||
print(f"✓ {metric} metric: Score = {results[0].score:.4f}")
|
||||
|
||||
await store.close()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Run all example."""
|
||||
print("\n" + "=" * 60)
|
||||
print("MongoDBStore Comprehensive Test Suite")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
# await example_basic_operations()
|
||||
# await example_filter_search()
|
||||
# await example_multiple_chunks()
|
||||
await example_distance_metrics()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✓ All tests completed successfully!")
|
||||
print("=" * 60)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n✗ Test failed with error: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
164
examples/functionality/vector_store/oceanbase/README.md
Normal file
164
examples/functionality/vector_store/oceanbase/README.md
Normal file
@@ -0,0 +1,164 @@
|
||||
# OceanBase Vector Store
|
||||
|
||||
This example demonstrates how to use **OceanBaseStore** for vector storage and semantic search in AgentScope.
|
||||
It includes CRUD operations, metadata filtering, document chunking, and distance metric tests.
|
||||
|
||||
### Quick Start
|
||||
|
||||
Install dependencies (including `pyobvector`):
|
||||
|
||||
```bash
|
||||
pip install -e .[full]
|
||||
```
|
||||
|
||||
Start seekdb (a minimal OceanBase-compatible instance):
|
||||
|
||||
```bash
|
||||
docker run -d -p 2881:2881 oceanbase/seekdb
|
||||
```
|
||||
|
||||
Run the example script:
|
||||
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
|
||||
> **Note:** The script defaults to `127.0.0.1:2881`, user `root`, database `test`.
|
||||
> If you use a multi-tenant OceanBase account (e.g., `root@test`), override via environment variables.
|
||||
|
||||
## Usage
|
||||
|
||||
### Initialize Store
|
||||
|
||||
```python
|
||||
from agentscope.rag import OceanBaseStore
|
||||
|
||||
store = OceanBaseStore(
|
||||
collection_name="test_collection",
|
||||
dimensions=768,
|
||||
distance="COSINE",
|
||||
uri="127.0.0.1:2881",
|
||||
user="root",
|
||||
password="",
|
||||
db_name="test",
|
||||
)
|
||||
```
|
||||
|
||||
### Add Documents
|
||||
|
||||
```python
|
||||
from agentscope.rag import Document, DocMetadata
|
||||
from agentscope.message import TextBlock
|
||||
|
||||
doc = Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(type="text", text="Your document text"),
|
||||
doc_id="doc_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3],
|
||||
)
|
||||
|
||||
await store.add([doc])
|
||||
```
|
||||
|
||||
### Search
|
||||
|
||||
```python
|
||||
results = await store.search(
|
||||
query_embedding=[0.1, 0.2, 0.3],
|
||||
limit=5,
|
||||
score_threshold=0.9,
|
||||
)
|
||||
```
|
||||
|
||||
### Filter Search
|
||||
|
||||
```python
|
||||
client = store.get_client()
|
||||
table = client.load_table(collection_name="test_collection")
|
||||
|
||||
results = await store.search(
|
||||
query_embedding=[0.1, 0.2, 0.3],
|
||||
limit=5,
|
||||
flter=[table.c["doc_id"].like("doc%")],
|
||||
)
|
||||
```
|
||||
|
||||
> Note: The parameter name is `flter` (missing the "i") to avoid clashing with
|
||||
> Python's built-in `filter` and follows the underlying library's convention.
|
||||
|
||||
### Delete
|
||||
|
||||
```python
|
||||
client = store.get_client()
|
||||
table = client.load_table(collection_name="test_collection")
|
||||
|
||||
await store.delete(where=[table.c["doc_id"] == "doc_1"])
|
||||
```
|
||||
|
||||
## Distance Metrics
|
||||
|
||||
| Metric | Description | Best For |
|
||||
|--------|-------------|----------|
|
||||
| **COSINE** | Cosine similarity | Text embeddings (recommended) |
|
||||
| **L2** | Euclidean distance | Spatial data |
|
||||
| **IP** | Inner product | Recommendation systems |
|
||||
|
||||
## Filter Expressions
|
||||
|
||||
Build filters using SQLAlchemy expressions and pass them via `flter`:
|
||||
|
||||
```python
|
||||
table = store.get_client().load_table("test_collection")
|
||||
|
||||
filters = [
|
||||
table.c["doc_id"] == "doc_1",
|
||||
table.c["doc_id"].like("prefix%"),
|
||||
table.c["chunk_id"] >= 0,
|
||||
]
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
### Access Underlying Client
|
||||
|
||||
```python
|
||||
client = store.get_client()
|
||||
stats = client.get_collection_stats(collection_name="test_collection")
|
||||
```
|
||||
|
||||
### Document Metadata
|
||||
|
||||
- `content`: Text content (TextBlock)
|
||||
- `doc_id`: Unique document identifier
|
||||
- `chunk_id`: Chunk position (0-indexed)
|
||||
- `total_chunks`: Total chunks in document
|
||||
|
||||
## FAQ
|
||||
|
||||
**What embedding dimension should I use?**
|
||||
Match your embedding model's output dimension (e.g., 768 for BERT, 1536 for OpenAI ada-002).
|
||||
|
||||
**Can I change the distance metric after creation?**
|
||||
No, create a new collection with the desired metric.
|
||||
|
||||
**How do I clean up test data?**
|
||||
Drop the collection via the underlying client or remove the seekdb container volume.
|
||||
|
||||
## Environment Variables
|
||||
|
||||
The script supports the following environment variables to override connection settings:
|
||||
|
||||
```bash
|
||||
export OCEANBASE_URI="127.0.0.1:2881"
|
||||
export OCEANBASE_USER="root"
|
||||
export OCEANBASE_PASSWORD=""
|
||||
export OCEANBASE_DB="test"
|
||||
```
|
||||
|
||||
## References
|
||||
|
||||
- [OceanBase Vector Store](https://github.com/oceanbase/pyobvector)
|
||||
- [AgentScope RAG Tutorial](https://doc.agentscope.io/tutorial/task_rag.html)
|
||||
350
examples/functionality/vector_store/oceanbase/main.py
Normal file
350
examples/functionality/vector_store/oceanbase/main.py
Normal file
@@ -0,0 +1,350 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Example of using OceanBaseStore in AgentScope RAG system."""
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from agentscope.rag import (
|
||||
OceanBaseStore,
|
||||
Document,
|
||||
DocMetadata,
|
||||
)
|
||||
from agentscope.message import TextBlock
|
||||
|
||||
|
||||
def _create_store(
|
||||
collection_name: str,
|
||||
dimensions: int = 4,
|
||||
distance: str = "COSINE",
|
||||
) -> OceanBaseStore:
|
||||
return OceanBaseStore(
|
||||
collection_name=collection_name,
|
||||
dimensions=dimensions,
|
||||
distance=distance,
|
||||
uri=os.getenv("OCEANBASE_URI", "127.0.0.1:2881"),
|
||||
user=os.getenv("OCEANBASE_USER", "root"),
|
||||
password=os.getenv("OCEANBASE_PASSWORD", ""),
|
||||
db_name=os.getenv("OCEANBASE_DB", "test"),
|
||||
)
|
||||
|
||||
|
||||
async def example_basic_operations() -> None:
|
||||
"""The example of basic CRUD operations with OceanBaseStore."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 1: Basic CRUD Operations")
|
||||
print("=" * 60)
|
||||
|
||||
store = _create_store(collection_name="ob_basic_collection")
|
||||
store.get_client().drop_collection("ob_basic_collection")
|
||||
|
||||
print("✓ OceanBaseStore initialized")
|
||||
|
||||
test_docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Artificial Intelligence is the future",
|
||||
),
|
||||
doc_id="doc_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Machine Learning is a subset of AI",
|
||||
),
|
||||
doc_id="doc_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Deep Learning uses neural networks",
|
||||
),
|
||||
doc_id="doc_3",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
]
|
||||
|
||||
await store.add(test_docs)
|
||||
print(f"✓ Added {len(test_docs)} documents to the store")
|
||||
|
||||
query_embedding = [0.15, 0.25, 0.35, 0.45]
|
||||
results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=2,
|
||||
)
|
||||
|
||||
print(f"\n✓ Search completed, found {len(results)} results:")
|
||||
for i, result in enumerate(results, 1):
|
||||
print(f" {i}. Score: {result.score:.4f}")
|
||||
print(f" Content: {result.metadata.content}")
|
||||
print(f" Doc ID: {result.metadata.doc_id}")
|
||||
|
||||
results_filtered = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=5,
|
||||
score_threshold=0.9,
|
||||
)
|
||||
print(f"\n✓ Search with threshold (>0.9): {len(results_filtered)} results")
|
||||
|
||||
client = store.get_client()
|
||||
table = client.load_table(collection_name="ob_basic_collection")
|
||||
await store.delete(where=[table.c["doc_id"] == "doc_2"])
|
||||
print("\n✓ Deleted document with doc_id='doc_2'")
|
||||
|
||||
results_after_delete = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=5,
|
||||
)
|
||||
print(f"✓ After deletion: {len(results_after_delete)} documents remain")
|
||||
|
||||
print(f"\n✓ Got MilvusLikeClient: {type(client).__name__}")
|
||||
|
||||
|
||||
async def example_filter_search() -> None:
|
||||
"""The example of search with metadata filtering."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 2: Search with Metadata Filtering")
|
||||
print("=" * 60)
|
||||
|
||||
store = _create_store(collection_name="ob_filter_collection")
|
||||
client = store.get_client()
|
||||
client.drop_collection("ob_filter_collection")
|
||||
|
||||
docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Python is a programming language",
|
||||
),
|
||||
doc_id="prog_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Java is used for enterprise applications",
|
||||
),
|
||||
doc_id="prog_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Neural networks are used in AI",
|
||||
),
|
||||
doc_id="ai_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Deep learning requires GPUs",
|
||||
),
|
||||
doc_id="ai_2",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.4, 0.5, 0.6, 0.7],
|
||||
),
|
||||
]
|
||||
|
||||
await store.add(docs)
|
||||
print(f"✓ Added {len(docs)} documents with different doc_id prefixes")
|
||||
|
||||
query_embedding = [0.25, 0.35, 0.45, 0.55]
|
||||
all_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
)
|
||||
print(f"\n✓ Search without filter: {len(all_results)} results")
|
||||
for i, result in enumerate(all_results, 1):
|
||||
print(
|
||||
f" {i}. Doc ID: {result.metadata.doc_id}, "
|
||||
f"Score: {result.score:.4f}",
|
||||
)
|
||||
|
||||
table = client.load_table(collection_name="ob_filter_collection")
|
||||
prog_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
flter=[table.c["doc_id"].like("prog%")],
|
||||
)
|
||||
print("\n✓ Search with filter (doc_id like 'prog%'):")
|
||||
for i, result in enumerate(prog_results, 1):
|
||||
print(
|
||||
f" {i}. Doc ID: {result.metadata.doc_id}, "
|
||||
f"Score: {result.score:.4f}",
|
||||
)
|
||||
|
||||
ai_results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=4,
|
||||
flter=[table.c["doc_id"].like("ai%")],
|
||||
)
|
||||
print("\n✓ Search with filter (doc_id like 'ai%'):")
|
||||
for i, result in enumerate(ai_results, 1):
|
||||
print(
|
||||
f" {i}. Doc ID: {result.metadata.doc_id}, "
|
||||
f"Score: {result.score:.4f}",
|
||||
)
|
||||
|
||||
|
||||
async def example_multiple_chunks() -> None:
|
||||
"""The example of documents with multiple chunks."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 3: Documents with Multiple Chunks")
|
||||
print("=" * 60)
|
||||
|
||||
store = _create_store(collection_name="ob_chunks_collection")
|
||||
store.get_client().drop_collection("ob_chunks_collection")
|
||||
|
||||
chunks = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Chapter 1: Introduction to AI",
|
||||
),
|
||||
doc_id="book_1",
|
||||
chunk_id=0,
|
||||
total_chunks=3,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Chapter 2: Machine Learning Basics",
|
||||
),
|
||||
doc_id="book_1",
|
||||
chunk_id=1,
|
||||
total_chunks=3,
|
||||
),
|
||||
embedding=[0.2, 0.3, 0.4, 0.5],
|
||||
),
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text="Chapter 3: Deep Learning Advanced",
|
||||
),
|
||||
doc_id="book_1",
|
||||
chunk_id=2,
|
||||
total_chunks=3,
|
||||
),
|
||||
embedding=[0.3, 0.4, 0.5, 0.6],
|
||||
),
|
||||
]
|
||||
|
||||
await store.add(chunks)
|
||||
print(f"✓ Added document with {len(chunks)} chunks")
|
||||
|
||||
query_embedding = [0.2, 0.3, 0.4, 0.5]
|
||||
results = await store.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=3,
|
||||
)
|
||||
|
||||
print("\n✓ Search results for multi-chunk document:")
|
||||
for i, result in enumerate(results, 1):
|
||||
chunk_info = (
|
||||
f"{result.metadata.chunk_id}/{result.metadata.total_chunks}"
|
||||
)
|
||||
print(f" {i}. Chunk {chunk_info}")
|
||||
print(f" Content: {result.metadata.content}")
|
||||
print(f" Score: {result.score:.4f}")
|
||||
|
||||
|
||||
async def example_distance_metrics() -> None:
|
||||
"""The example of different distance metrics."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 4: Different Distance Metrics")
|
||||
print("=" * 60)
|
||||
|
||||
metrics = ["COSINE", "L2", "IP"]
|
||||
|
||||
for metric in metrics:
|
||||
print(f"\n--- Testing {metric} metric ---")
|
||||
collection_name = f"ob_{metric}_collection"
|
||||
store = _create_store(
|
||||
collection_name=collection_name,
|
||||
distance=metric,
|
||||
)
|
||||
store.get_client().drop_collection(collection_name)
|
||||
|
||||
docs = [
|
||||
Document(
|
||||
metadata=DocMetadata(
|
||||
content=TextBlock(
|
||||
type="text",
|
||||
text=f"Test doc for {metric}",
|
||||
),
|
||||
doc_id=f"doc_{metric}_1",
|
||||
chunk_id=0,
|
||||
total_chunks=1,
|
||||
),
|
||||
embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
),
|
||||
]
|
||||
|
||||
await store.add(docs)
|
||||
results = await store.search(
|
||||
query_embedding=[0.1, 0.2, 0.3, 0.4],
|
||||
limit=1,
|
||||
)
|
||||
|
||||
print(f"✓ {metric} metric: Score = {results[0].score:.4f}")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Run all example."""
|
||||
print("\n" + "=" * 60)
|
||||
print("OceanBaseStore Comprehensive Test Suite")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
await example_basic_operations()
|
||||
await example_filter_search()
|
||||
await example_multiple_chunks()
|
||||
await example_distance_metrics()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✓ All tests completed successfully!")
|
||||
print("=" * 60)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n✗ Test failed with error: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user