ReMe/test/cookbook/appworld/appworld_react_agent.py
jinliyl d0c9d89092
feat(memory): add ContextChecker component for context size management (#144)
* feat(memory): add ContextChecker component for context size management

* refactor(memory): restructure file-based memory tools and update imports

* docs(readme): update documentation with detailed architecture and components

* docs(readme): update Chinese documentation with enhanced memory management diagrams

* refactor(cookbook): move cookbook files to test directory and clean up docs

* docs(readme): update link path for old version documentation

* docs(readme): update documentation with improved architecture diagrams and component details

* docs(readme): update documentation with improved clarity and structure

* refactor(docs): update in-memory memory documentation

* docs(readme): add experiment reproduction link to quickstart guide
2026-03-06 23:43:42 +08:00

350 lines
13 KiB
Python

# flake8: noqa: E402, E501
import os
from typing import List, Any
from tqdm import tqdm
os.environ["APPWORLD_ROOT"] = "."
from dotenv import load_dotenv
load_dotenv("../../.env")
import re
import time
import json
import ray
import requests
import datetime
from appworld import AppWorld, load_task_ids
from jinja2 import Template
from loguru import logger
from openai import OpenAI
from prompt import NEW_PROMPT_TEMPLATE
@ray.remote
class AppworldReactAgent:
"""A minimal ReAct Agent for AppWorld tasks."""
def __init__(
self,
index: int,
task_ids: List[str],
experiment_name: str,
model_name: str = "qwen3-8b",
temperature: float = 0.9,
max_interactions: int = 30,
max_response_size: int = 129024,
num_trials: int = 1,
use_memory: bool = False,
use_memory_addition: bool = False,
use_memory_deletion: bool = False,
delete_freq: int = 10,
freq_threshold: int = 5,
utility_threshold: float = 0.5,
memory_base_url: str = "http://0.0.0.0:8002/",
memory_workspace_id: str = "appworld_v1",
):
self.index: int = index
self.task_ids: List[str] = task_ids
self.experiment_name: str = experiment_name
self.model_name: str = model_name
self.temperature: float = temperature
self.max_interactions: int = max_interactions
self.max_response_size: int = max_response_size
self.num_trials: int = num_trials
self.use_memory: bool = use_memory
self.use_memory_addition: bool = use_memory_addition if use_memory else False
self.use_memory_deletion: bool = use_memory_deletion if use_memory else False
self.delete_freq: int = delete_freq
self.freq_threshold: int = freq_threshold
self.utility_threshold: float = utility_threshold
self.memory_base_url: str = memory_base_url
self.memory_workspace_id: str = memory_workspace_id
self.llm_client = OpenAI()
self.history: List[List[List[dict]]] = [[] for _ in range(num_trials)]
self.retrieved_memory_list: List[List[List[Any]]] = [[] for _ in range(num_trials)]
for run_id in range(num_trials):
for _ in range(len(task_ids)):
self.retrieved_memory_list[run_id].append([])
self.history[run_id].append([])
def call_llm(self, messages: list) -> str:
for i in range(100):
try:
response = self.llm_client.chat.completions.create(
model=self.model_name,
messages=messages,
temperature=self.temperature,
extra_body={"enable_thinking": False},
seed=0,
)
return response.choices[0].message.content
except Exception as e:
logger.exception(f"encounter error with {e.args}")
time.sleep(1 + i * 10)
return "call llm error"
def prompt_messages(self, run_id, task_index, previous_memories: None, world: AppWorld):
app_descriptions = json.dumps(
[{"name": k, "description": v} for (k, v) in world.task.app_descriptions.items()],
indent=1,
)
dictionary = {"supervisor": world.task.supervisor, "app_descriptions": app_descriptions}
sys_prompt = Template(NEW_PROMPT_TEMPLATE.lstrip()).render(dictionary)
query = world.task.instruction
if self.use_memory:
if len(previous_memories) == 0:
response = self.get_memory(world.task.instruction)
if response and "memory_list" in response["metadata"]:
self.retrieved_memory_list[run_id][task_index] = response["metadata"]["memory_list"]
task_memory = response["answer"]
logger.info(f"loaded task_memory: {task_memory}")
query = (
"Task:\n"
+ query
+ "\n\nSome Related Experience to help you to complete the task:\n"
+ re.sub(r"(?i)\bMemory\s*(\d+)\s*[:]", r"Experience \1:", task_memory)
)
else:
formatted_memories = []
for i, memory in enumerate(previous_memories, 1):
condition = memory["when_to_use"]
memory_content = memory["content"]
memory_text = f"Experience {i}:\n When to use: {condition}\n Content: {memory_content}\n"
formatted_memories.append(memory_text)
query = (
"Task:\n"
+ query
+ "\n\nSome Related Experience to help you to complete the task:\n"
+ "\n".join(formatted_memories)
)
messages = [
{"role": "system", "content": sys_prompt},
{"role": "user", "content": query},
]
self.history[run_id][task_index] = messages
@staticmethod
def get_reward(world) -> float:
tracker = world.evaluate()
num_passes = len(tracker.passes)
num_failures = len(tracker.failures)
return num_passes / (num_passes + num_failures)
def extract_code_and_fix_content(
self,
text: str,
ignore_multiple_calls=True,
) -> tuple[str, str]:
full_code_regex = r"```python\n(.*?)```"
partial_code_regex = r".*```python\n(.*)"
original_text = text
output_code = ""
match_end = 0
# Handle multiple calls
for re_match in re.finditer(full_code_regex, original_text, flags=re.DOTALL):
code = re_match.group(1).strip()
if ignore_multiple_calls:
text = original_text[: re_match.end()]
return code, text
output_code += code + "\n"
match_end = re_match.end()
# check for partial code match at end (no terminating ```) following the last match
partial_match = re.match(
partial_code_regex,
original_text[match_end:],
flags=re.DOTALL,
)
if partial_match:
output_code += partial_match.group(1).strip()
# terminated due to stop condition. Add stop condition to output.
if not text.endswith("\n"):
text = text + "\n"
text = text + "```"
if len(output_code) == 0:
return text, text
else:
return output_code, text
def execute(self):
result = []
counter = 0
for task_index, task_id in enumerate(tqdm(self.task_ids, desc=f"ray_index={self.index}")):
t_result = None
previous_memories = []
# Run each task num_trials times
for run_id in range(self.num_trials):
start_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
with AppWorld(task_id=task_id, experiment_name=f"{self.experiment_name}_run_{run_id}") as world:
before_score = self.get_reward(world)
for i in range(self.max_interactions):
if i == 0:
self.prompt_messages(
run_id=run_id,
task_index=task_index,
previous_memories=previous_memories,
world=world,
)
code_msg = self.call_llm(self.history[run_id][task_index])
code, text = self.extract_code_and_fix_content(code_msg)
self.history[run_id][task_index].append({"role": "assistant", "content": code})
output = world.execute(code)
# if len(output) > self.max_response_size:
# # logger.warning(f"output exceed max size={len(output)}")
# output = output[: self.max_response_size]
self.history[run_id][task_index].append(
{"role": "user", "content": "Output:\n```\n" + output + "```\n\n"},
)
if world.task_completed():
break
after_score = self.get_reward(world)
uplift_score = after_score - before_score
if self.use_memory:
if self.use_memory_addition:
new_traj_list = [
self.get_traj_from_task_history(task_id, self.history[run_id][task_index], after_score),
]
previous_memories = self.add_memory(new_traj_list)
if after_score != 1:
self.delete_memory_by_ids([mem["memory_id"] for mem in previous_memories])
# update the freq & utility attributes of retrieved memories
update_utility: bool = after_score == 1
self.update_memory_information(self.retrieved_memory_list[run_id][task_index], update_utility)
counter += 1
if self.use_memory_deletion: # and counter % self.delete_freq == 0:
self.delete_memory()
t_result = {
"task_id": world.task_id,
"run_id": run_id,
"experiment_name": self.experiment_name,
"task_completed": world.task_completed(),
"before_score": before_score,
"after_score": after_score,
"uplift_score": uplift_score,
"task_history": self.history[run_id][task_index],
"task_start_time": start_time,
}
if after_score == 1:
break
result.append(t_result)
return result
def handle_api_response(self, response: requests.Response):
"""Handle API response with proper error checking"""
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.text)
return None
return response.json()
def get_memory(self, query: str):
"""Retrieve relevant task memories based on a query"""
response = requests.post(
url=f"{self.memory_base_url}retrieve_task_memory",
json={
"workspace_id": self.memory_workspace_id,
"query": query,
},
)
result = self.handle_api_response(response)
if not result:
return None
logger.info(f"query: {query}, response: {result}")
return result
def get_traj_from_task_history(self, task_id: str, task_history: list, reward: float):
pattern = r"\n\nSome Related Experience to help you to complete the task:.*"
task_history[1]["content"] = re.sub(pattern, "", task_history[1]["content"], flags=re.DOTALL)
return {
"task_id": task_id,
"messages": task_history,
"score": reward,
}
def add_memory(self, trajectories):
"""Generate a summary of conversation messages and create task memories"""
response = requests.post(
url=f"{self.memory_base_url}summary_task_memory",
json={
"workspace_id": self.memory_workspace_id,
"trajectories": trajectories,
},
)
result = self.handle_api_response(response)
if not result:
return []
# Extract memory list from response
memory_list = result.get("metadata", {}).get("memory_list", [])
print(f"Task memory list created: {len(memory_list)} memories")
return memory_list
def delete_memory_by_ids(self, memory_ids):
response = requests.post(
url=f"{self.memory_base_url}vector_store",
json={
"workspace_id": self.memory_workspace_id,
"action": "delete_ids",
"memory_ids": memory_ids,
},
)
response.raise_for_status()
def update_memory_information(self, memory_list, update_utility: bool = False):
response = requests.post(
url=f"{self.memory_base_url}record_task_memory",
json={
"workspace_id": self.memory_workspace_id,
"memory_dicts": memory_list,
"update_utility": update_utility,
},
)
response.raise_for_status()
logger.info(response.json())
def delete_memory(self):
response = requests.post(
url=f"{self.memory_base_url}delete_task_memory",
json={
"workspace_id": self.memory_workspace_id,
"freq_threshold": self.freq_threshold,
"utility_threshold": self.utility_threshold,
},
)
response.raise_for_status()
def main():
dataset_name = "train"
task_ids = load_task_ids(dataset_name)
agent = AppworldReactAgent(index=0, task_ids=task_ids[0:1], experiment_name=dataset_name, num_trials=1)
result = agent.execute()
logger.info(f"result={json.dumps(result)}")
if __name__ == "__main__":
main()