From a2b7a616647d0ce977d8cd01243eb7ad79c8d989 Mon Sep 17 00:00:00 2001
From: "jinli.yl" <jinli.yl@alibaba-inc.com>
Date: Thu, 27 Nov 2025 16:50:37 +0800
Subject: [PATCH] feat(working_memory): add working memory demo with ReAct
 agent

---
 .../react_agent_with_working_memory.py        | 137 +++++++++++
 cookbook/working_memory/work_memory_demo.py   | 117 ++++++++++
 docs/cookbook/working/quick_start.md          | 220 ++++++++++++++++++
 .../task/memory_validation_prompt.yaml        |   4 +-
 4 files changed, 476 insertions(+), 2 deletions(-)
 create mode 100644 cookbook/working_memory/react_agent_with_working_memory.py
 create mode 100644 cookbook/working_memory/work_memory_demo.py
 create mode 100644 docs/cookbook/working/quick_start.md

diff --git a/cookbook/working_memory/react_agent_with_working_memory.py b/cookbook/working_memory/react_agent_with_working_memory.py
new file mode 100644
index 00000000..52e3fdc6
--- /dev/null
+++ b/cookbook/working_memory/react_agent_with_working_memory.py
@@ -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
diff --git a/cookbook/working_memory/work_memory_demo.py b/cookbook/working_memory/work_memory_demo.py
new file mode 100644
index 00000000..612a637a
--- /dev/null
+++ b/cookbook/working_memory/work_memory_demo.py
@@ -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())
diff --git a/docs/cookbook/working/quick_start.md b/docs/cookbook/working/quick_start.md
new file mode 100644
index 00000000..a4bf6893
--- /dev/null
+++ b/docs/cookbook/working/quick_start.md
@@ -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
diff --git a/reme_ai/summary/task/memory_validation_prompt.yaml b/reme_ai/summary/task/memory_validation_prompt.yaml
index 2738d50d..44287e28 100644
--- a/reme_ai/summary/task/memory_validation_prompt.yaml
+++ b/reme_ai/summary/task/memory_validation_prompt.yaml
@@ -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?