Skip to content

Commit aec81ff

Browse files
Rasic2Rasic2
committed
refactor：重构状态管理逻辑以简化当前步骤处理和验证
1 parent 0334053 commit aec81ff

18 files changed

Lines changed: 298 additions & 190 deletions

File tree

‎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: 5 additions & 3 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 CURRENT_STEP
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,

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,8 @@
6464
)
6565
from agents.matmaster_agent.state import (
6666
CURRENT_STEP,
67+
CURRENT_STEP_DESCRIPTION,
6768
RECOMMEND_PARAMS,
68-
STEP_DESCRIPTION,
6969
)
7070
from agents.matmaster_agent.sub_agents.tools import ALL_TOOLS
7171
from agents.matmaster_agent.utils.event_utils import (
@@ -215,7 +215,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
215215
)
216216

217217
self.tool_call_info_agent.instruction = gen_tool_call_info_instruction(
218-
user_prompt=current_step[STEP_DESCRIPTION],
218+
user_prompt=current_step[CURRENT_STEP_DESCRIPTION],
219219
agent_prompt=self.instruction,
220220
tool_doc=tool_doc,
221221
tool_schema=tool_schema,

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
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
3029
from agents.matmaster_agent.state import CURRENT_STEP
3130
from agents.matmaster_agent.utils.event_utils import (
@@ -143,7 +142,7 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
143142
# Only Query Job Result
144143
step_title = ctx.session.state.get('step_title', {}).get(
145144
'title',
146-
f"{i18n.t(ctx.session.state['separate_card_info'])} {ctx.session.state['plan_index'] + 1}: {current_step_tool_name}",
145+
current_step_tool_name,
147146
)
148147
for matmaster_flow_event in context_function_event(
149148
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: 21 additions & 16 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 CURRENT_STEP
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')
@@ -270,10 +275,10 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
270275
ctx, self.name, '工具参数无变化，本次跳过执行', ModelRole
271276
):
272277
yield _info_event
273-
update_plan = copy.deepcopy(ctx.session.state['plan'])
274-
update_plan['steps'][ctx.session.state['plan_index']][
275-
'status'
276-
] = PlanStepStatusEnum.FAILED
277-
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+
)
278283
else:
279284
yield event

0 commit comments

Comments
 (0)