add dr_agents planner memory · InternScience/InternAgent@fa8c3ee · GitHub
Skip to content

Commit fa8c3ee

Browse files
noc-turnenoc-turne
committed
add dr_agents planner memory
1 parent f6ef9bb commit fa8c3ee

4 files changed

Lines changed: 130 additions & 20 deletions

File tree

‎docs/memory_module.md‎

Lines changed: 2 additions & 1 deletion

‎internagent/mas/agents/dr_agents/agent_kb/README.md‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,10 @@ experience knowledge base and retrieved before executing similar future subtasks
1313

1414
## How it is wired into `dr_agents`
1515

16-
- **Retrieval** — `agents/task/execution_agent.py` calls `call_hybrid_search('base'/'append', task)`
16+
- **Planner-side retrieval** — `agents/global_planner_agent.py` retrieves aligned planning
17+
experience and injects it into the planner prompt (`GLOBAL_PLANNER_MEMORY_BLOCK`) to guide
18+
graph construction. An external `set_memory_prompt(...)` hook is also provided.
19+
- **Executor-side retrieval** — `agents/task/execution_agent.py` calls `call_hybrid_search('base'/'append', task)`
1720
before a subtask and injects the results into `EXECUTION_PROMPT_WITHMEMORY`.
1821
- **Write-back** — `workflow/task.py` collects the subtask trace, distills it with
1922
`MEMORY_REASONING_PROMPT`, and calls `call_appendkb(...)` after the task finishes.

‎internagent/mas/agents/dr_agents/agents/global_planner_agent.py‎

Lines changed: 105 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,26 @@
1818

1919
from datetime import datetime
2020

21+
# SPM (Strategy-Procedural Memory, paper Sec. 2.4.1) — planner-side strategic priors.
22+
# The same experience KB used by the executor also guides the planner toward globally
23+
# coherent reasoning graphs. Imported lazily/safely so the planner still runs when the
24+
# agent_kb service/package is unavailable.
25+
try:
26+
from agent_kb.example import call_hybrid_search
27+
except Exception: # pragma: no cover - agent_kb is optional
28+
call_hybrid_search = None
29+
30+
31+
def _is_truthy(val) -> bool:
32+
"""True for boolean True or the strings 'true'/'True' (env-var friendly)."""
33+
return val is True or (isinstance(val, str) and val.strip().lower() == "true")
34+
35+
36+
def _need_memory_enabled(config) -> bool:
37+
"""Resolve the SPM on/off switch: env `need_memory` takes precedence, then config."""
38+
cfg = config or {}
39+
return _is_truthy(os.getenv("need_memory", cfg.get("need_memory", False)))
40+
2141

2242
class GlobalPlannerAgent(BaseAgent):
2343
"""
@@ -47,10 +67,60 @@ def __init__(self, model: str = "deepseek-r1", tools: List[Dict[str, Any]] = Non
4767
self.config,
4868
default_name="GLOBAL_PLANNER_PROMPT"
4969
)
50-
70+
71+
# SPM planner-side: strategic priors block, appended to the base prompt only when
72+
# SPM is active. Kept separate so the default planner prompt is unchanged when off.
73+
planner_memory_block_config = {
74+
"prompt_path": self.config.get("planner_memory_block_path"),
75+
"prompt_name": self.config.get("planner_memory_block_name")
76+
}
77+
self.planner_memory_block = load_prompt(
78+
planner_memory_block_config,
79+
default_name="GLOBAL_PLANNER_MEMORY_BLOCK"
80+
)
81+
self.prompt_template_with_memory = self.prompt_template + "\n" + self.planner_memory_block
82+
self.need_memory = _need_memory_enabled(self.config)
83+
self.memory_prompt = None # external override hook (set_memory_prompt)
84+
self.base_retrieve_score_threshold = self.config.get('base_retrieve_score_threshold', 0.3)
85+
self.append_retrieve_score_threshold = self.config.get('append_retrieve_score_threshold', 0.65)
86+
5187
self.max_iter = self.config.get('max_iter', 5)
5288
self.max_retries = self.config.get('max_retries', 3)
5389
self.max_nodes = self.config.get('max_nodes', 8)
90+
91+
def set_memory_prompt(self, memory_prompt: str):
92+
"""外部注入 SPM 记忆提示（覆盖内部检索）。paper Sec. 2.4.1 的 planner-side 钩子。"""
93+
self.memory_prompt = memory_prompt
94+
95+
def _retrieve_planner_memory(self, question: str):
96+
"""
97+
从 agent_kb 检索与当前问题相似的历史 planning 经验（strategic priors），
98+
组装成注入 planner prompt 的 memory_prompt。检索不可用或失败时返回 None（不影响规划）。
99+
"""
100+
if not (self.need_memory and call_hybrid_search is not None):
101+
if self.need_memory and call_hybrid_search is None:
102+
self.logger.warning("need_memory 已开启，但 agent_kb 客户端不可用，planner 不使用记忆")
103+
return None
104+
try:
105+
base_similar = call_hybrid_search('base', question, top_k=2, weights={"text": 0.3, "semantic": 0.7})
106+
append_similar = call_hybrid_search('append', question, top_k=2, weights={"text": 0.3, "semantic": 0.7})
107+
refs = []
108+
for exp in (base_similar or []):
109+
if exp.get('total_score', 0) > self.base_retrieve_score_threshold:
110+
refs.append((exp.get('query'), exp.get('agent_experience', '') or ''))
111+
for exp in (append_similar or []):
112+
if exp.get('total_score', 0) > self.append_retrieve_score_threshold:
113+
refs.append((exp.get('query'), exp.get('agent_experience', '') or ''))
114+
if not refs:
115+
return None
116+
parts = []
117+
for q, planning_exp in refs:
118+
parts.append(f"[Similar past task]\n{q}\n[Planning experience]\n{planning_exp}")
119+
self.logger.info(f"SPM planner 检索到 {len(refs)} 条 planning 经验")
120+
return "\n\n".join(parts)
121+
except Exception as e:
122+
self.logger.warning(f"SPM planner 检索失败，本次规划不使用记忆: {e}")
123+
return None
54124

55125
def build_graph_from_plan(self, plan_result: Dict[str, Any]) -> DirectedGraph:
56126
"""
@@ -123,34 +193,47 @@ def mark_node_executed(self, node_id: str) -> bool:
123193

124194
return self.graph.set_node_status(node_id, NodeExecutionStatus.EXECUTED)
125195

126-
def execute_one_step(self, input_data: str, question: str, current_iter = 1, max_iter = 5, additional_info = None, tools = None) -> Dict[str, Any]:
196+
def execute_one_step(self, input_data: str, question: str, current_iter = 1, max_iter = 5, additional_info = None, tools = None, memory_prompt = None) -> Dict[str, Any]:
127197
"""
128198
执行任务规划，将自然语言任务拆解为子任务并构建依赖关系图
129-
199+
130200
Args:
131201
input_data: 中间图，包含nodes和edges
132202
question: 问题描述
133203
current_iter: 当前迭代次数
134204
max_iter: 最大迭代次数
135-
205+
memory_prompt: SPM 检索到的 planning 先验（可选）；提供时使用带记忆的 prompt
206+
136207
Returns:
137208
包含 nodes 和 edges 的字典，表示任务分解结果
138209
"""
139210

140-
formatted_prompt = self.prompt_template.format(
141-
graph=input_data,
142-
question=question,
143-
max_iter=str(max_iter),
144-
current_iter=str(current_iter),
145-
additional_info=additional_info,
146-
tools=json.dumps(tools),
147-
max_nodes=str(self.max_nodes)
148-
)
149-
211+
if memory_prompt:
212+
formatted_prompt = self.prompt_template_with_memory.format(
213+
graph=input_data,
214+
question=question,
215+
max_iter=str(max_iter),
216+
current_iter=str(current_iter),
217+
additional_info=additional_info,
218+
tools=json.dumps(tools),
219+
max_nodes=str(self.max_nodes),
220+
memory_prompt=memory_prompt
221+
)
222+
else:
223+
formatted_prompt = self.prompt_template.format(
224+
graph=input_data,
225+
question=question,
226+
max_iter=str(max_iter),
227+
current_iter=str(current_iter),
228+
additional_info=additional_info,
229+
tools=json.dumps(tools),
230+
max_nodes=str(self.max_nodes)
231+
)
232+
150233
# 调用模型生成响应
151234
response = self._generate_plan(formatted_prompt)
152-
153-
return response
235+
236+
return response
154237

155238
def execute(self, input_data: str, file_path = None, additional_info = None) -> Dict[str, Any]:
156239
"""
@@ -182,17 +265,21 @@ def execute(self, input_data: str, file_path = None, additional_info = None) ->
182265
tool_info = self.tool_manager.get_simple_tool_info(tool)
183266
tool_info_list.append(tool_info)
184267

268+
# SPM planner-side: 取一次 planning 先验（外部注入优先，否则从 agent_kb 检索），
269+
# 整个规划循环复用，避免每轮重复检索。
270+
memory_prompt = self.memory_prompt or self._retrieve_planner_memory(input_data)
271+
185272
while True: # 检查生成的图中是否存在环，如果存在环，则重新构建图
186273
idx = 1
187274
while idx <= max_iter:
188275
# 添加重试机制
189276
max_retries = self.max_retries
190277
retry_count = 0
191278
response = None
192-
279+
193280
while retry_count < max_retries:
194281
try:
195-
response = self.execute_one_step(json.dumps(graph), input_data, current_iter=idx, max_iter=max_iter, additional_info=additional_info, tools=tool_info_list)
282+
response = self.execute_one_step(json.dumps(graph), input_data, current_iter=idx, max_iter=max_iter, additional_info=additional_info, tools=tool_info_list, memory_prompt=memory_prompt)
196283
if response is not None:
197284
break # 成功获得响应，跳出重试循环
198285
except Exception as e:

‎internagent/mas/agents/dr_agents/prompts/default_prompts.py‎

Lines changed: 19 additions & 0 deletions

0 commit comments

Comments
 (0)