Skip to content

Commit 652b2c8

Browse files
committed
refactor:简化执行和计划确认逻辑以支持更动态的步骤生成与执行
1 parent d27abb9 commit 652b2c8

28 files changed

Lines changed: 597 additions & 742 deletions

File tree

agents/matmaster_agent/callback.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,13 @@
1717
from agents.matmaster_agent.model import UserContent
1818
from agents.matmaster_agent.prompt import get_user_content_lang
1919
from agents.matmaster_agent.services.quota import check_quota_service, use_quota_service
20-
from agents.matmaster_agent.state import ERROR_DETAIL, ERROR_OCCURRED, PLAN, UPLOAD_FILE
20+
from agents.matmaster_agent.state import (
21+
ERROR_DETAIL,
22+
ERROR_OCCURRED,
23+
HISTORY_STEPS,
24+
PLAN,
25+
UPLOAD_FILE,
26+
)
2127
from agents.matmaster_agent.utils.helper_func import get_user_id
2228

2329
logger = logging.getLogger(__name__)
@@ -131,6 +137,10 @@ async def matmaster_prepare_state(
131137
callback_context.state['separate_card_info'] = callback_context.state.get(
132138
'separate_card_info', ''
133139
)
140+
# 历史的所有执行过程
141+
callback_context.state[HISTORY_STEPS] = callback_context.state.get(
142+
HISTORY_STEPS, []
143+
)
134144

135145

136146
async def matmaster_set_lang(

agents/matmaster_agent/constant.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
MATMASTER_AGENT_NAME = 'matmaster_agent'
1717

1818
ModelRole = 'model'
19+
UserRole = 'user'
1920

2021
Transfer2Agent = 'transfer_to_agent'
2122

agents/matmaster_agent/core_agents/base_agents/mcp_agent.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@
5050
store_tool_result_in_memory,
5151
)
5252
from agents.matmaster_agent.model import CostFuncType
53-
from agents.matmaster_agent.state import PLAN
53+
from agents.matmaster_agent.state import CURRENT_STEP, CURRENT_STEP_RESULT
5454
from agents.matmaster_agent.style import tool_response_failed_card
5555
from agents.matmaster_agent.utils.event_utils import (
5656
all_text_event,
@@ -240,8 +240,10 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
240240
raise
241241

242242
parsed_tool_result = await parse_result(ctx, dict_result)
243-
logger.info(
244-
f'{ctx.session.id} parsed_tool_result = {parsed_tool_result}'
243+
post_execution_step = copy.deepcopy(ctx.session.state[CURRENT_STEP])
244+
post_execution_step[CURRENT_STEP_RESULT] = parsed_tool_result
245+
yield update_state_event(
246+
ctx, state_delta={CURRENT_STEP: post_execution_step}
245247
)
246248
for _frontend_render_event in frontend_render_event(
247249
ctx,
@@ -263,9 +265,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
263265
not event.partial
264266
and event.content.parts[0].text
265267
== 'All Function Calls Are Occurred Before, Continue'
266-
and ctx.session.state[PLAN]['steps'][
267-
ctx.session.state['plan_index']
268-
]['status']
268+
and ctx.session.state[CURRENT_STEP]['status']
269269
== PlanStepStatusEnum.PROCESS
270270
):
271271
for _info_event in all_text_event(

agents/matmaster_agent/core_agents/base_agents/schema_agent.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from agents.matmaster_agent.core_agents.comp_agents.dntransfer_climit_agent import (
1414
CombinedDisallowTransferAndContentLimitMixin,
1515
)
16+
from agents.matmaster_agent.logger import PrefixFilter
1617
from agents.matmaster_agent.utils.event_utils import (
1718
context_function_event,
1819
is_function_call,
@@ -22,6 +23,8 @@
2223
from agents.matmaster_agent.utils.helper_func import extract_json_from_string
2324

2425
logger = logging.getLogger(__name__)
26+
logger.addFilter(PrefixFilter(MATMASTER_AGENT_NAME))
27+
logger.setLevel(logging.INFO)
2528

2629

2730
class SchemaAgent(ErrorHandleLlmAgent):
@@ -41,7 +44,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
4144
r',(\s*[}\]])', r'\1', repaired_raw_text
4245
) # 移除尾随逗号
4346
logger.info(
44-
f'[{MATMASTER_AGENT_NAME}]:[{ctx.session.id}] repaired_raw_text = {repaired_raw_text}'
47+
f'{ctx.session.id} {self.name} repaired_raw_text = {repaired_raw_text}'
4548
)
4649
schema_info: dict = json.loads(repaired_raw_text)
4750

agents/matmaster_agent/core_agents/comp_agents/recommend_summary_agent/agent.py

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -55,15 +55,18 @@
5555
)
5656
from agents.matmaster_agent.flow_agents.model import PlanStepStatusEnum
5757
from agents.matmaster_agent.llm_config import MatMasterLlmConfig
58-
from agents.matmaster_agent.locales import i18n
5958
from agents.matmaster_agent.logger import PrefixFilter
6059
from agents.matmaster_agent.model import ToolCallInfoSchema
6160
from agents.matmaster_agent.prompt import (
6261
GLOBAL_INSTRUCTION,
6362
GLOBAL_SCHEMA_INSTRUCTION,
6463
get_vocabulary_enforce_prompt,
6564
)
66-
from agents.matmaster_agent.state import RECOMMEND_PARAMS, STEP_DESCRIPTION
65+
from agents.matmaster_agent.state import (
66+
CURRENT_STEP,
67+
CURRENT_STEP_DESCRIPTION,
68+
RECOMMEND_PARAMS,
69+
)
6770
from agents.matmaster_agent.sub_agents.tools import ALL_TOOLS
6871
from agents.matmaster_agent.utils.event_utils import (
6972
context_function_event,
@@ -178,9 +181,7 @@ def summary_agent(self) -> DisallowTransferAndContentLimitLlmAgent:
178181
@override
179182
async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, None]:
180183
# 根据计划来
181-
current_step = ctx.session.state['plan']['steps'][
182-
ctx.session.state['plan_index']
183-
]
184+
current_step = ctx.session.state[CURRENT_STEP]
184185
current_step_tool_name = current_step['tool_name']
185186

186187
# 连接 tool-server,获取doc和函数声明
@@ -214,7 +215,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
214215
)
215216

216217
self.tool_call_info_agent.instruction = gen_tool_call_info_instruction(
217-
user_prompt=current_step[STEP_DESCRIPTION],
218+
user_prompt=current_step[CURRENT_STEP_DESCRIPTION],
218219
agent_prompt=self.instruction,
219220
tool_doc=tool_doc,
220221
tool_schema=tool_schema,
@@ -318,7 +319,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
318319

319320
step_title = ctx.session.state.get('step_title', {}).get(
320321
'title',
321-
f"{i18n.t(ctx.session.state['separate_card_info'])} {ctx.session.state['plan_index'] + 1}: {current_step_tool_name}",
322+
current_step_tool_name,
322323
)
323324
for matmaster_flow_event in context_function_event(
324325
ctx,

agents/matmaster_agent/core_agents/public_agents/job_agents/agent.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,8 @@
2525
SubmitRenderAgent,
2626
)
2727
from agents.matmaster_agent.flow_agents.model import PlanStepStatusEnum
28-
from agents.matmaster_agent.locales import i18n
2928
from agents.matmaster_agent.logger import PrefixFilter
29+
from agents.matmaster_agent.state import CURRENT_STEP
3030
from agents.matmaster_agent.utils.event_utils import (
3131
all_text_event,
3232
context_function_event,
@@ -132,9 +132,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
132132
async for result_event in self.result_agent.run_async(ctx):
133133
yield result_event
134134

135-
current_step = ctx.session.state['plan']['steps'][
136-
ctx.session.state['plan_index']
137-
]
135+
current_step = ctx.session.state[CURRENT_STEP]
138136
current_step_tool_name = current_step['tool_name']
139137
current_step_status = current_step['status']
140138
if current_step_status in [
@@ -144,7 +142,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
144142
# Only Query Job Result
145143
step_title = ctx.session.state.get('step_title', {}).get(
146144
'title',
147-
f"{i18n.t(ctx.session.state['separate_card_info'])} {ctx.session.state['plan_index'] + 1}: {current_step_tool_name}",
145+
current_step_tool_name,
148146
)
149147
for matmaster_flow_event in context_function_event(
150148
ctx,

agents/matmaster_agent/core_agents/public_agents/job_agents/result_core_agent/agent.py

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,14 @@
2525
from agents.matmaster_agent.core_agents.public_agents.job_agents.result_core_agent.prompt import (
2626
ResultCoreAgentDescription,
2727
)
28+
from agents.matmaster_agent.flow_agents.step_utils import get_current_step
2829
from agents.matmaster_agent.logger import PrefixFilter
2930
from agents.matmaster_agent.services.job import (
3031
get_job_detail,
3132
parse_and_prepare_err,
3233
parse_and_prepare_results,
3334
)
35+
from agents.matmaster_agent.state import CURRENT_STEP, CURRENT_STEP_STATUS
3436
from agents.matmaster_agent.utils.event_utils import (
3537
all_text_event,
3638
context_function_event,
@@ -115,14 +117,8 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
115117
if status != 'Running':
116118
# 更新状态
117119
plan_status = 'success' if status == 'Finished' else 'failed'
118-
update_plan = copy.deepcopy(ctx.session.state['plan'])
119-
logger.info(
120-
f'{ctx.session.id} plan_index = {ctx.session.state['plan_index']}'
121-
)
122-
update_plan['steps'][ctx.session.state['plan_index']][
123-
'status'
124-
] = plan_status
125-
120+
post_execution_step = copy.deepcopy(get_current_step(ctx))
121+
post_execution_step[CURRENT_STEP_STATUS] = plan_status
126122
update_long_running_jobs = copy.deepcopy(
127123
ctx.session.state['long_running_jobs']
128124
)
@@ -131,7 +127,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
131127
ctx,
132128
state_delta={
133129
'long_running_jobs': update_long_running_jobs,
134-
'plan': update_plan,
130+
CURRENT_STEP: post_execution_step,
135131
},
136132
)
137133

agents/matmaster_agent/core_agents/public_agents/job_agents/submit_core_agent/agent.py

Lines changed: 22 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,11 @@
1515
DisallowTransferAndContentLimitMCPAgent,
1616
)
1717
from agents.matmaster_agent.flow_agents.model import PlanStepStatusEnum
18+
from agents.matmaster_agent.flow_agents.step_utils import get_current_step
1819
from agents.matmaster_agent.locales import i18n
1920
from agents.matmaster_agent.logger import PrefixFilter
2021
from agents.matmaster_agent.model import BohrJobInfo, DFlowJobInfo
21-
from agents.matmaster_agent.state import PLAN
22+
from agents.matmaster_agent.state import CURRENT_STEP, CURRENT_STEP_STATUS
2223
from agents.matmaster_agent.style import tool_response_failed_card
2324
from agents.matmaster_agent.utils.event_utils import (
2425
all_text_event,
@@ -99,11 +100,13 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
99100
yield tool_response_failed_event
100101

101102
# 更新 plan 为失败
102-
update_plan = copy.deepcopy(ctx.session.state['plan'])
103-
update_plan['steps'][ctx.session.state['plan_index']][
104-
'status'
105-
] = 'failed'
106-
yield update_state_event(ctx, state_delta={'plan': update_plan})
103+
post_execution_step = copy.deepcopy(get_current_step(ctx))
104+
post_execution_step[CURRENT_STEP_STATUS] = (
105+
PlanStepStatusEnum.FAILED
106+
)
107+
yield update_state_event(
108+
ctx, state_delta={CURRENT_STEP: post_execution_step}
109+
)
107110

108111
raise RuntimeError('Tool Execution Failed')
109112
dict_result = load_tool_response(first_part)
@@ -189,12 +192,14 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
189192
yield tool_response_failed_event
190193

191194
# 更新 plan 为失败
192-
update_plan = copy.deepcopy(ctx.session.state['plan'])
193-
update_plan['steps'][ctx.session.state['plan_index']][
194-
'status'
195-
] = 'failed'
195+
post_execution_step = copy.deepcopy(
196+
get_current_step(ctx)
197+
)
198+
post_execution_step[CURRENT_STEP_STATUS] = (
199+
PlanStepStatusEnum.FAILED
200+
)
196201
yield update_state_event(
197-
ctx, state_delta={'plan': update_plan}
202+
ctx, state_delta={CURRENT_STEP: post_execution_step}
198203
)
199204

200205
raise RuntimeError('Tool Execution Failed')
@@ -263,19 +268,17 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
263268
not event.partial
264269
and event.content.parts[0].text
265270
== 'All Function Calls Are Occurred Before, Continue'
266-
and ctx.session.state[PLAN]['steps'][
267-
ctx.session.state['plan_index']
268-
]['status']
271+
and ctx.session.state[CURRENT_STEP]['status']
269272
== PlanStepStatusEnum.PROCESS
270273
):
271274
for _info_event in all_text_event(
272275
ctx, self.name, '工具参数无变化,本次跳过执行', ModelRole
273276
):
274277
yield _info_event
275-
update_plan = copy.deepcopy(ctx.session.state['plan'])
276-
update_plan['steps'][ctx.session.state['plan_index']][
277-
'status'
278-
] = PlanStepStatusEnum.FAILED
279-
yield update_state_event(ctx, state_delta={'plan': update_plan})
278+
post_execution_step = copy.deepcopy(ctx.session.state[CURRENT_STEP])
279+
post_execution_step['status'] = PlanStepStatusEnum.FAILED
280+
yield update_state_event(
281+
ctx, state_delta={CURRENT_STEP: post_execution_step}
282+
)
280283
else:
281284
yield event

0 commit comments

Comments
 (0)