scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-08-28 05:25:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 11

feat(working_memory): add working memory demo with ReAct agent

Browse source
This commit is contained in:
jinli.yl 2025-11-27 16:50:37 +08:00
parent 0ee91e912c
commit a2b7a61664
4 changed files with 476 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

137
cookbook/working_memory/react_agent_with_working_memory.py Normal file
View file

@ -0,0 +1,137 @@
import json
from typing import List, Dict
from flowllm.core.enumeration import Role
from flowllm.core.schema import ToolCall, Message
from flowllm.core.token import BaseToken
from flowllm.core.utils import load_env
from loguru import logger
load_env()
from flowllm.core.llm import OpenAICompatibleLLM
from flowllm.core.utils import FastMcpClient, HttpClient
class ReactAgent:
"""
A simple ReAct-style agent that:
- Talks to an OpenAI-compatible LLM backend.
- Uses ReMe's working memory summary flow (`summary_working_memory`) to
automatically compress context before each reasoning step.
- Calls two MCP tools exposed by the ReMe backend:
- `grep_working_memory`: search working memory by keyword / regex.
- `read_working_memory`: read a specific segment of working memory.
This demo focuses on how to integrate ReMe working memory into a ReAct loop,
rather than on complex agent logic.
"""
def __init__(self,
model_name="",
max_steps: int = 50):
# You can replace this with your own LLM wrapper if needed.
self.llm = OpenAICompatibleLLM(model_name=model_name)
self.max_steps = max_steps
@staticmethod
def token_count(messages: List[Message]):
"""
Count tokens for a list of `Message` objects using FlowLLM's token counter.
This is useful for measuring compression effects of working memory summary.
"""
return BaseToken().token_count(messages)
async def run(self, messages: List[Message]):
"""
Run the ReAct loop with ReMe working memory.
Requirements (services to start before running this demo):
- ReMe MCP server, exposing `grep_working_memory` and `read_working_memory` tools, e.g.:
`reme backend=mcp mcp.port=8002 ...`
- ReMe HTTP server, exposing the `summary_working_memory` flow, e.g.:
`reme backend=http http.port=8003 ...`
The loop roughly does:
1. Summarize / compress the conversation with `summary_working_memory`.
2. Ask the LLM what to do next (possibly call tools).
3. Execute tool calls via MCP if requested, append tool results.
4. Repeat until the model no longer requests tools or `max_steps` is reached.
"""
# Prepare all available tools from the MCP server.
tool_dict: Dict[str, ToolCall] = {}
async with FastMcpClient("reme_mcp_server", {
"type": "sse",
"url": "http://0.0.0.0:8002/sse",
}) as mcp_client, HttpClient(base_url="http://localhost:8003") as http_client:
tool_calls = await mcp_client.list_tool_calls()
for tool_call in tool_calls:
if tool_call.name in ["grep_working_memory", "read_working_memory"]:
tool_dict[tool_call.name] = tool_call
# Log the tool call schema in Qwen3-compatible format for debugging.
# (This is the standard "tool" format for Qwen3 / BaiLian.)
tool_call_str = json.dumps(tool_call.simple_input_dump(), ensure_ascii=False, indent=2)
logger.info(f"tool_call {tool_call.name} {tool_call_str}")
# Main ReAct loop.
for i in range(self.max_steps):
# Before every LLM call, run `summary_working_memory` to:
# - compress long histories,
# - offload detailed context into working memory storage,
# - keep the recent message(s) for short-term reasoning.
result = await http_client.execute_flow("summary_working_memory",
messages=[x.simple_dump() for x in messages],
working_summary_mode="auto",
compact_ratio_threshold=0.75,
max_total_tokens=20000,
max_tool_message_tokens=2000,
group_token_threshold=None,
keep_recent_count=1,
store_dir="./test_working_memory")
# Convert the API result back into `Message` objects for the LLM.
messages = [Message(**x) for x in result.answer]
# Ask the LLM what to do next.
# You can plug in your own tool-calling strategy here.
assistant_message: Message = await self.llm.achat(messages=messages, tools=[
tool_dict["grep_working_memory"],
tool_dict["read_working_memory"],
])
messages.append(assistant_message)
if not assistant_message.tool_calls:
# If the LLM does not request any tools, we assume it has finished.
break
for j, tool_call in enumerate(assistant_message.tool_calls):
if tool_call.name not in tool_dict:
logger.exception(f"unknown tool_call.name={tool_call.name}")
continue
logger.info(f"round{i + 1}.{j} submit tool_calls={tool_call.name} "
f"argument={tool_call.argument_dict}")
# Execute the tool via MCP and parse the result.
result = await mcp_client.call_tool(tool_call.name,
arguments=tool_call.argument_dict,
parse_result=True)
# Attach the tool result as a TOOL-role message so the LLM
# can see and reason about it in the next step.
messages.append(Message(
role=Role.TOOL,
tool_call_id=tool_call.id,
content=result,
))
return messages

117
cookbook/working_memory/work_memory_demo.py Normal file
View file

@ -0,0 +1,117 @@
import asyncio
from flowllm.core.enumeration import Role
from flowllm.core.schema import ToolCall, Message
from loguru import logger
from react_agent_with_working_memory import ReactAgent
async def main():
"""
End-to-end demo for using `ReactAgent` with ReMe working memory.
The scenario is:
- The user asks about the ReMe project's README.
- The assistant (ReAct agent) should first use a `web_search` tool call
(simulated here) to obtain README content.
- The full README (multiplied by 4 to simulate a long context) is injected
as a TOOL message.
- Then the user asks a question that requires reading that long content:
"According to the README, what is the quantitative effect of task memory
in appworld?"
- `ReactAgent` plus ReMe working memory compresses this long context and
tries to answer while keeping token usage low.
"""
# A fake tool_call_id to tie TOOL messages back to the ASSISTANT's tool call.
tool_call_id = "call_6596dafa2a6a46f7a217da"
# Load the ReMe README as the "raw data" we want the agent to read.
# In a real scenario, this could be code, docs, logs, etc.
with open("../../README.md", encoding="utf-8") as f:
readme_content = f.read()
# Build a conversation that simulates:
# 1. System message with instructions (how to use Grep/ReadFile tools).
# 2. User requesting to search README.
# 3. Assistant triggering a `web_search` tool call (simulated).
# 4. TOOL message that returns the README content (repeated 4 times).
# 5. A follow-up user question that requires understanding that README.
messages = [
Message(
role=Role.SYSTEM,
content=(
"You are a helpful assistant. "
"请先使用`Grep`匹配关键词或者正则表达式所在行数，然后通过`ReadFile`读取位置附近的代码。"
"如果没有找到匹配项，永远不要放弃尝试，尝试其他的参数，比如只搜索部分关键词。"
"`Grep`之后通过 `ReadFile` 命令，你可以从指定偏移位置`offset`+长度`limit`开始查看内容，不要超过100行。"
"如果当前内容不足，`ReadFile` 命令也可以不断尝试不同的`offset`和`limit`参数"
),
),
Message(
role=Role.USER,
content="搜索下reme项目的的README内容",
),
Message(
role=Role.ASSISTANT,
content="",
tool_calls=[
ToolCall(
**{
"index": 0,
"id": tool_call_id,
"function": {
"arguments": '{"query": "readme"}',
"name": "web_search",
},
"type": "function",
},
),
],
),
# Simulate the tool result: the README content is returned as if from `web_search`.
# We repeat it 4 times to create a large context and better showcase the
# compression ability of ReMe working memory.
Message(
role=Role.TOOL,
content=readme_content * 4,
tool_call_id=tool_call_id,
),
Message(
role=Role.USER,
content="根据readme回答task memory在appworld的效果是多少，需要具体的数值",
),
]
# Measure token count before running the agent, so we can see compression ratio.
origin_token_count = ReactAgent.token_count(messages)
# You can change this to any OpenAI-compatible model name that your backend exposes.
# For example:
# model_name = "qwen3-30b-a3b-instruct-2507"
model_name = "qwen3-coder-30b-a3b-instruct"
# Initialize the agent with a maximum number of reasoning/tool steps.
agent = ReactAgent(model_name=model_name, max_steps=50)
# Run the ReAct loop with ReMe working memory summarization.
messages = await agent.run(messages)
# Count tokens again to see how much the context has been compressed.
after_token_count = ReactAgent.token_count(messages)
logger.info(f"result: {messages}")
# Example numbers (from a previous run):
# origin_token_count: 24586 after_token_count: 1565 compress_ratio=0.06
logger.info(
f"origin_token_count: {origin_token_count} "
f"after_token_count: {after_token_count} "
f"compress_ratio={after_token_count / origin_token_count:.2f}"
)
if __name__ == "__main__":
# Run the async demo entrypoint.
asyncio.run(main())

220
docs/cookbook/working/quick_start.md Normal file
View file

@ -0,0 +1,220 @@
# Working Memory Demo
This demo showcases how to use ReMe's working memory capabilities with a ReAct agent. The working memory system automatically manages context by compressing and summarizing conversation history, enabling efficient long-context processing.
## Installation
### Install from PyPI (Recommended)
```bash
pip install reme-ai
```
### Install from Source
```bash
git clone https://github.com/agentscope-ai/ReMe.git
cd ReMe
pip install .
```
### Environment Configuration
Copy `example.env` to `.env` and modify the corresponding parameters:
```bash
FLOW_LLM_API_KEY=sk-xxxx
FLOW_LLM_BASE_URL=https://xxxx/v1
FLOW_EMBEDDING_API_KEY=sk-xxxx
FLOW_EMBEDDING_BASE_URL=https://xxxx/v1
```
## Starting the Services
Before running the demo, you need to start both the HTTP and MCP services:
### Start MCP Service
```bash
reme backend=mcp mcp.port=8002
```
The MCP service provides tools for working memory management including:
- `grep_working_memory`: Search for content in working memory
- `read_working_memory`: Read specific sections of working memory
### Start HTTP Service
```bash
reme backend=http http.port=8003
```
The HTTP service provides the flow execution endpoint for memory operations.
## Running the Demo
Once both services are running, execute the demo:
```bash
cd cookbook/working_memory
python work_memory_demo.py
```
### What the Demo Does
The demo simulates a scenario where:
1. A large README content is loaded (repeated 4 times to create a long context)
2. The agent needs to search through this content and extract specific information
3. Working memory automatically compresses the context from ~24,586 tokens to ~1,565 tokens (compression ratio: 0.06)
4. The agent can still accurately answer questions about the content
## Core Code Explanation
### ReactAgent with Working Memory (`react_agent_with_working_memory.py`)
#### 1. Agent Initialization
```python
class ReactAgent:
def __init__(self, model_name="", max_steps: int = 50):
# Use your own LLM class
self.llm = OpenAICompatibleLLM(model_name=model_name)
self.max_steps = max_steps
```
The agent is initialized with an LLM model and a maximum number of reasoning steps.
#### 2. Service Connection
```python
async with FastMcpClient("reme_mcp_server", {
"type": "sse",
"url": "http://0.0.0.0:8002/sse",
}) as mcp_client, HttpClient(base_url="http://localhost:8003") as http_client:
```
The agent connects to both:
- **MCP Client**: For tool execution (grep, read operations)
- **HTTP Client**: For flow execution (memory summarization)
#### 3. Tool Registration
```python
tool_calls = await mcp_client.list_tool_calls()
for tool_call in tool_calls:
if tool_call.name in ["grep_working_memory", "read_working_memory"]:
tool_dict[tool_call.name] = tool_call
```
The agent registers working memory tools that will be available to the LLM.
> Note: `summary_working_memory` is **not** an MCP tool.
> It is a **flow** exposed by the HTTP service and is invoked via `HttpClient.execute_flow`,
> as shown in the next section.
#### 4. Working Memory Summarization (Key Feature)
```python
result = await http_client.execute_flow("summary_working_memory",
messages=[x.simple_dump() for x in messages],
working_summary_mode="auto",
compact_ratio_threshold=0.75,
max_total_tokens=20000,
max_tool_message_tokens=2000,
group_token_threshold=None,
keep_recent_count=1,
store_dir="./test_working_memory")
messages = [Message(**x) for x in result.answer]
```
**This is the core of working memory management.** Before each LLM call:
- **`working_summary_mode="auto"`**: Automatically decides when to compress
- **`compact_ratio_threshold=0.75`**: Triggers compression when context exceeds 75% of max tokens
- **`max_total_tokens=20000`**: Maximum total tokens allowed
- **`max_tool_message_tokens=2000`**: Maximum tokens per tool message
- **`keep_recent_count=1`**: Keeps the most recent message uncompressed
- **`store_dir`**: Directory to store compressed memory
The summarization process:
1. Analyzes the current message history
2. Identifies compressible content (especially long tool outputs)
3. Compresses/summarizes old messages while preserving semantic information
4. Returns a condensed message list that maintains context
#### 5. ReAct Loop
```python
for i in range(self.max_steps):
# Summarize working memory before each LLM call
result = await http_client.execute_flow("summary_working_memory", ...)
messages = [Message(**x) for x in result.answer]
# LLM generates next action
assistant_message = await self.llm.achat(messages=messages, tools=[...])
messages.append(assistant_message)
if not assistant_message.tool_calls:
break
# Execute tools
for tool_call in assistant_message.tool_calls:
result = await mcp_client.call_tool(tool_call.name,
arguments=tool_call.argument_dict)
messages.append(Message(role=Role.TOOL, content=result, ...))
```
The ReAct loop:
1. **Compress**: Summarize working memory to reduce context size
2. **Reason**: LLM decides what tool to use
3. **Act**: Execute the tool
4. **Observe**: Add tool result to messages
5. Repeat until task is complete or max steps reached
### Benefits of Working Memory
1. **Context Efficiency**: Reduces token usage by ~94% (24,586 → 1,565 tokens in the demo)
2. **Cost Reduction**: Lower token counts mean lower API costs
3. **Performance**: Faster inference with smaller contexts
4. **Scalability**: Handle much longer conversations and tool outputs
5. **Accuracy**: Maintains semantic information despite compression
## Model Configuration
The demo uses an OpenAI-compatible LLM configured via environment variables:
- **`FLOW_LLM_API_KEY` / `FLOW_LLM_BASE_URL`**: LLM API credentials and endpoint
- The model name is specified in `work_memory_demo.py`, for example:
```python
model_name = "qwen3-coder-30b-a3b-instruct"
agent = ReactAgent(model_name=model_name, max_steps=50)
```
You can change `model_name` to any model that your backend supports, as long as it follows the OpenAI-compatible API.
## Expected Output
When running the demo, you should see:
- Token count before compression: ~24,586 tokens
- Token count after compression: ~1,565 tokens
- Compression ratio: ~0.06 (6% of original size)
- The agent successfully answers the question about task memory performance in AppWorld
## Customization
You can customize the working memory behavior by adjusting parameters in the `summary_working_memory` call:
- **`compact_ratio_threshold`**: Lower values trigger compression earlier
- **`max_total_tokens`**: Adjust based on your model's context window
- **`max_tool_message_tokens`**: Control individual tool output size
- **`keep_recent_count`**: Keep more recent messages uncompressed for better context
## Troubleshooting
1. **Services not starting**: Ensure ports 8002 and 8003 are available
2. **Connection errors**: Verify both MCP and HTTP services are running
3. **API errors**: Check your `.env` file has valid API keys and endpoints
4. **Memory errors**: Adjust `max_total_tokens` based on your available memory

4
reme_ai/summary/task/memory_validation_prompt.yaml
View file

@ -1,8 +1,8 @@
task_memory_validation_prompt: |
You are an expert AI analyst tasked with validating the quality and usefulness of extracted step-level task memories.
Your task is to access whether the extracted task memory is actionable, accurate, and valuable for future agent executions.
VALIDATION CRITERIA:
● ACTIONABILITY: Is the task memory specific enough to guide future actions?
● ACCURACY: Does the task memory correctly reflect the patterns observed?