Skip to content

Commit e0c6017

Browse files
fix: save history in state variale with Event update
1 parent 89f7529 commit e0c6017

1 file changed

Lines changed: 17 additions & 6 deletions

File tree

MaxKernel/auto_agent/subagents/pipeline_agent.py

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ async def _run_async_impl(
7979
self._clear_iteration_metrics(ctx)
8080

8181
if self._should_end_at_step(ctx, iteration, "plan"):
82+
yield self._create_history_event(ctx)
8283
iteration += 1
8384
continue
8485

@@ -87,6 +88,7 @@ async def _run_async_impl(
8788
async for event in self.implement_agent.run_async(ctx):
8889
yield event
8990
if self._should_end_at_step(ctx, iteration, "implement"):
91+
yield self._create_history_event(ctx)
9092
iteration += 1
9193
continue
9294

@@ -106,10 +108,12 @@ async def _run_async_impl(
106108
self._save_iteration_files_and_snapshot(
107109
ctx, iteration, step_name="validate"
108110
)
111+
yield self._create_history_event(ctx)
109112
iteration += 1
110113
continue
111114

112115
if self._should_end_at_step(ctx, iteration, "validate"):
116+
yield self._create_history_event(ctx)
113117
iteration += 1
114118
continue
115119

@@ -146,10 +150,12 @@ async def _run_async_impl(
146150
self._save_iteration_files_and_snapshot(
147151
ctx, iteration, step_name="test_run"
148152
)
153+
yield self._create_history_event(ctx)
149154
iteration += 1
150155
continue
151156

152157
if self._should_end_at_step(ctx, iteration, "test_run"):
158+
yield self._create_history_event(ctx)
153159
iteration += 1
154160
continue
155161

@@ -158,6 +164,7 @@ async def _run_async_impl(
158164
async for event in self.autotune_agent.run_async(ctx):
159165
yield event
160166
if self._should_end_at_step(ctx, iteration, "autotune"):
167+
yield self._create_history_event(ctx)
161168
iteration += 1
162169
continue
163170

@@ -166,17 +173,13 @@ async def _run_async_impl(
166173
async for event in self.profile_agent.run_async(ctx):
167174
yield event
168175
if self._should_end_at_step(ctx, iteration, "profile"):
176+
yield self._create_history_event(ctx)
169177
iteration += 1
170178
continue
171179

172180
self._save_iteration_files_and_snapshot(ctx, iteration)
173181

174-
yield Event(
175-
author=self.name,
176-
actions=EventActions(
177-
state_delta={"history": ctx.session.state.get("history", [])}
178-
),
179-
)
182+
yield self._create_history_event(ctx)
180183

181184
# Step 7: Check if improvement is needed
182185
# needs_improvement = ctx.session.state.get("needs_improvement", False)
@@ -210,6 +213,14 @@ async def _run_async_impl(
210213
),
211214
)
212215

216+
def _create_history_event(self, ctx: InvocationContext) -> Event:
217+
return Event(
218+
author=self.name,
219+
actions=EventActions(
220+
state_delta={"history": ctx.session.state.get("history", [])}
221+
),
222+
)
223+
213224
def _should_end_at_step(
214225
self, ctx: InvocationContext, iteration: int, step_name: str
215226
) -> bool:

0 commit comments

Comments
 (0)