From 8940de3eacb884ccc830b37798bf13b7cb25dd09 Mon Sep 17 00:00:00 2001 From: Taras Tarasiuk Date: Sun, 18 Oct 2026 22:46:07 +0200 Subject: [PATCH] feat: cross-iteration memory --- src/agentomics/agents/steps/base.py | 58 ++++++-- src/agentomics/agents/steps/iteration_plan.py | 14 ++ src/agentomics/runtime/iterations_lessons.py | 126 ++++++++++++++++++ src/agentomics/utils/config.py | 5 + .../test_step_iterations_lessons.py | 65 +++++++++ 5 files changed, 260 insertions(+), 8 deletions(-) create mode 100644 src/agentomics/runtime/iterations_lessons.py create mode 100644 tests/integration/test_step_iterations_lessons.py diff --git a/src/agentomics/agents/steps/base.py b/src/agentomics/agents/steps/base.py index dade12d2..65fd2f47 100644 --- a/src/agentomics/agents/steps/base.py +++ b/src/agentomics/agents/steps/base.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, Field from pydantic.json_schema import SkipJsonSchema from pydantic_ai import Agent, ModelRetry, RunContext +from pydantic_ai.messages import ModelMessage from pydantic_ai.models import Model from pydantic_ai.settings import ModelSettings @@ -22,6 +23,7 @@ from agentomics.runtime.conda_utils import export_shared_environment_descriptor from agentomics.runtime.filesystem import chown_tree_to_root, rewrite_symlinks_to_absolute from agentomics.runtime.git_checkpoints import commit_step_checkpoint +from agentomics.runtime.iterations_lessons import read_iterations_lessons, update_iterations_lessons from agentomics.runtime.read_write_utils import ( get_new_rundir_files, load_current_iteration_base_prompt, @@ -30,6 +32,7 @@ ) from agentomics.runtime.step_outputs import load_step_outputs, save_step_output from agentomics.utils.config import Config +from agentomics.utils.exceptions import IterationRunFailed from agentomics.utils.providers.provider import Provider @@ -54,7 +57,12 @@ async def run(self) -> None: self.on_step_success(self.build_simulated_output()) return self.on_step_start() - output = await self._execute() + try: + output = await self._execute() + except Exception as error: + await self.on_step_end(error=error) + raise + await self.on_step_end() self.on_step_success(output) def should_run(self) -> bool: @@ -79,6 +87,11 @@ def on_step_start(self) -> None: check=True, ) + async def on_step_end(self, error: Exception | None = None) -> None: + iteration = load_current_iteration_index(self.config) + status = "FAILED" if error else "SUCCEEDED" + _console.print(f"[bold purple]ITERATION {iteration} | {self.display_name} STEP {status}[/bold purple]") + def on_step_success(self, output: Any) -> None: save_step_output(self.config, self.step_id, output) export_shared_environment_descriptor(self.config) @@ -121,6 +134,16 @@ def on_step_start(self) -> None: super().on_step_start() self._setup_injected_scripts() + async def on_step_end(self, error: Exception | None = None) -> None: + await super().on_step_end(error) + await update_iterations_lessons( + config=self.config, + model=self.iteration_plan_model, + step_name=self.display_name, + messages=self._step_messages, + error=error, + ) + def injected_scripts(self) -> list[Path]: return [] @@ -185,8 +208,18 @@ def build_user_prompt(self) -> str: if self.injected_scripts(): prompt += "\nYour current working directory contains a helpers/ package of read-only modules you can import." prompt += "\nSummarized outputs from your previous steps are in previous messages." + prompt += self._build_iterations_lessons() return prompt + def _build_iterations_lessons(self) -> str: + iterations_lessons = read_iterations_lessons(self.config) + if not iterations_lessons: + return "" + return ( + "\nLessons established so far in this run. Apply the ones relevant to this step instead of re-discovering them.\n" + f"\n{iterations_lessons}\n" + ) + def build_deps(self, step_started_at: datetime) -> dict[str, Any]: return {"start_time": step_started_at} @@ -216,15 +249,24 @@ def get_message_history(self): return history async def _execute(self) -> AgenticStepOutput: + self._step_messages: list[ModelMessage] = [] step_started_at = datetime.now() message_history = self.get_message_history() + history_length = len(message_history or []) user_prompt = self.build_user_prompt() agent = self.create_agent() - _, output = await run_agent( - agent=agent, - user_prompt=user_prompt, - max_steps=self.config.max_steps, - message_history=message_history, - deps=self.build_deps(step_started_at), - ) + try: + run_messages, output = await run_agent( + agent=agent, + user_prompt=user_prompt, + max_steps=self.config.max_steps, + message_history=message_history, + deps=self.build_deps(step_started_at), + ) + except IterationRunFailed as error: + # history_length: drop the previous steps' context, keep only this step's messages + self._step_messages = error.context_messages[history_length:] + raise + # history_length: drop the previous steps' context, keep only this step's messages + self._step_messages = run_messages[history_length:] return self._populate_files_created(output, step_started_at) diff --git a/src/agentomics/agents/steps/iteration_plan.py b/src/agentomics/agents/steps/iteration_plan.py index 2b012aa0..976772fb 100644 --- a/src/agentomics/agents/steps/iteration_plan.py +++ b/src/agentomics/agents/steps/iteration_plan.py @@ -10,6 +10,7 @@ from agentomics.agents.steps.base import AgenticStep, AgenticStepOutput from agentomics.agents.steps.data_split import DataSplitStep from agentomics.agents.steps.validation_evaluation import ValidationEvaluationStep +from agentomics.runtime.iterations_lessons import read_iterations_lessons from agentomics.runtime.read_write_utils import ( get_archived_iterations, get_last_successful_iteration, @@ -103,6 +104,7 @@ def step_prompt(self) -> str: tools_info = ", ".join(get_tool_names(self.config, self.config.tool_ids)) best_iteration_model_info = self._build_best_iteration_model_info() supplementary_info = self._build_supplementary_info() + iterations_lessons = self._build_iterations_lessons() return f""" @@ -114,6 +116,7 @@ def step_prompt(self) -> str: {iteration_history_info} + {iterations_lessons} {f"{exploration_info}" if exploration_info else ""} The main goal of the run is to maximize the hidden test set generalization performance (main metric:{self.config.val_metric}) that will use the 'best iteration model'. {best_iteration_model_info} @@ -157,6 +160,17 @@ def step_prompt(self) -> str: def get_message_history(self): return None + def _build_iterations_lessons(self) -> str: + iterations_lessons = read_iterations_lessons(self.config) + if not iterations_lessons: + return "" + return f""" + Lessons established so far in this run, including failed steps and iterations. They are automatically appended to the prompt of every step of this iteration, so don't repeat them in your instructions, but take them into account when planning. + + {iterations_lessons} + + """ + def _build_iteration_history_info(self) -> str: sorted_iteration_to_records = { iteration: self._load_iteration_history_record(iteration) diff --git a/src/agentomics/runtime/iterations_lessons.py b/src/agentomics/runtime/iterations_lessons.py new file mode 100644 index 00000000..36db9f64 --- /dev/null +++ b/src/agentomics/runtime/iterations_lessons.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import traceback +from collections.abc import Callable +from typing import Any + +from pydantic_ai import Agent +from pydantic_ai.messages import ( + ModelMessage, + RetryPromptPart, + TextPart, + ToolCallPart, + ToolReturnPart, +) +from pydantic_ai.models import Model +from pydantic_ai.settings import ModelSettings + +from agentomics.agents.agent_utils import run_agent +from agentomics.utils.config import Config +from agentomics.utils.exceptions import IterationRunFailed + +ITERATIONS_LESSONS_AGENT_MAX_REQUESTS = 5 +NO_NEW_ITERATIONS_LESSONS = "NO NEW LESSONS" +MAX_PART_CHARS = 1500 +MAX_TRANSCRIPT_CHARS = 60000 + +# Parts not listed here (system prompt, user prompt, thinking) are left out of the transcript +# Only the text needed for lessons is taken from each listed part +_PART_CONTENT: dict[type, Callable[[Any], str]] = { + ToolCallPart: lambda part: f"{part.tool_name} {part.args_as_json_str()}", + ToolReturnPart: lambda part: f"{part.tool_name} {part.model_response_str()}", + RetryPromptPart: lambda part: part.model_response(), + TextPart: lambda part: part.content, +} + +ITERATIONS_LESSONS_SYSTEM_PROMPT = """ +You maintain iterations_lessons.md, the memory of mistakes made by an autonomous ML agent that runs the same sequence of steps over several iterations on one dataset. +After every step, successful or failed, you receive the current file and the transcript of that step. + +A lesson is a mistake the agent made and what to do instead, such as: +- an error or dead end and its root cause +- a validator complaint and how it was satisfied +- an environment problem, such as a missing package, a timeout, or a wrong path +- a workflow problem, such as repeating the same tool call or not finishing the step +Don't record findings about the data or the model unless they caused a mistake, good practices, plans, or generic ML advice. Never invent anything the transcript doesn't show. + +Write each lesson as an instruction for the next agent, under about 25 words: +- : . . +For example: +- Training: `xgboost` is not installed. Use `sklearn.ensemble.HistGradientBoostingClassifier` instead. +- Any step: `write_python` rejects paths outside the current step folder. Save new scripts there. +Keep the names the next agent needs, such as columns, packages, and functions. Leave out full error messages, absolute paths, script names, names of functions or variables the agent wrote itself, and the story of how the fix was found. + +Return the complete new file as plain markdown, without code fences or commentary: +- Keep every existing lesson unless the transcript shows it is wrong. In that case, correct or remove it. +- Write one lesson per root cause, not one per symptom. If several errors come from the same mistake, describe it once. +- If the file already states the same mistake in other words, merge the new detail into that bullet instead of adding a near-duplicate. +- Keep the file under 60 bullets. +- If the step made no new mistakes, return exactly: NO NEW LESSONS +""" + + +def read_iterations_lessons(config: Config) -> str: + if not config.iterations_lessons_path.exists(): + return "" + return config.iterations_lessons_path.read_text(encoding="utf-8").strip() + + +async def update_iterations_lessons( + config: Config, + model: Model, + step_name: str, + messages: list[ModelMessage], + error: Exception | None = None, +) -> None: + try: + _, iterations_lessons = await run_agent( + agent=_create_iterations_lessons_agent(config, model), + user_prompt=_build_user_prompt(config, step_name, messages, error), + max_steps=ITERATIONS_LESSONS_AGENT_MAX_REQUESTS, + message_history=None, + verbose=False, + ) + except Exception: + print(f"WARNING: {Config.ITERATIONS_LESSONS_FILENAME} was not updated after step {step_name}.") + return + # todo: test the approach with structured lessons (a DTO instead of a plain string) and compare after observability metrics are introduced + if iterations_lessons.strip() != NO_NEW_ITERATIONS_LESSONS: + config.iterations_lessons_path.write_text(f"{iterations_lessons.strip()}\n", encoding="utf-8") + + +def _build_user_prompt(config: Config, step_name: str, messages: list[ModelMessage], error: Exception | None) -> str: + result = "It succeeded." if error is None else f"It failed with:\n{_format_error(error)}" + return f""" +Current {Config.ITERATIONS_LESSONS_FILENAME}: +{read_iterations_lessons(config)} + +Step {step_name} just ended. Below is the transcript of its agent, for you to analyze: +{_format_transcript(messages)} + +{result} +""" + + +def _format_transcript(messages: list[ModelMessage]) -> str: + parts = [part for message in messages for part in message.parts if type(part) in _PART_CONTENT] + contents = [(type(part).__name__, _PART_CONTENT[type(part)](part)) for part in parts] + lines = [f"[{name}] {content[:MAX_PART_CHARS]}" for name, content in contents if content.strip()] + return "\n".join(lines)[-MAX_TRANSCRIPT_CHARS:] + + +def _format_error(error: Exception) -> str: + trace = ( + error.exception_trace if isinstance(error, IterationRunFailed) else "".join(traceback.format_exception(error)) + ) + return trace[-MAX_PART_CHARS:] + + +def _create_iterations_lessons_agent(config: Config, model: Model) -> Agent[None, str]: + return Agent( + model=model, + system_prompt=ITERATIONS_LESSONS_SYSTEM_PROMPT, + model_settings=ModelSettings(temperature=config.temperature), + output_type=str, + retries=config.max_validation_retries, + ) diff --git a/src/agentomics/utils/config.py b/src/agentomics/utils/config.py index c33fe762..dd45c28b 100644 --- a/src/agentomics/utils/config.py +++ b/src/agentomics/utils/config.py @@ -24,6 +24,7 @@ class Config: ITERATION_STATE_FILENAME: ClassVar[str] = "iteration_state.json" ENVIRONMENT_DESCRIPTOR_FILENAME: ClassVar[str] = "environment.yml" STEP_OUTPUT_FILENAME: ClassVar[str] = "output.json" + ITERATIONS_LESSONS_FILENAME: ClassVar[str] = "iterations_lessons.md" DEFAULT_ITERATIONS: ClassVar[int] = 5 DEFAULT_SPLIT_ALLOWED_ITERATIONS: ClassVar[int] = 1 @@ -129,6 +130,10 @@ def dataset_dir(self) -> Path: def shared_dir(self) -> Path: return self.run_dir / self.SHARED_DIRNAME + @property + def iterations_lessons_path(self) -> Path: + return self.shared_dir / self.ITERATIONS_LESSONS_FILENAME + @property def config_path(self) -> Path: return self.shared_dir / self.CONFIG_FILENAME diff --git a/tests/integration/test_step_iterations_lessons.py b/tests/integration/test_step_iterations_lessons.py new file mode 100644 index 00000000..4e18c079 --- /dev/null +++ b/tests/integration/test_step_iterations_lessons.py @@ -0,0 +1,65 @@ +import asyncio +from unittest.mock import Mock + +import pytest +from pydantic_ai.messages import ModelResponse, TextPart +from pydantic_ai.models.function import FunctionModel + +from agentomics.agents.steps.data_representation import ( + DataRepresentationOutput, + DataRepresentationStep, +) +from agentomics.utils.exceptions import IterationRunFailed +from tests.helpers import create_model_calling_output_tool + +MOCK_ITERATIONS_LESSONS_RESPONSE = "- Representation: use one-hot encoding. All sequences have the same length." + + +def _failing_model() -> FunctionModel: + def respond(_messages, _agent_info): + raise RuntimeError("provider returned finish_reason='error'") + + return FunctionModel(respond) + + +def _iterations_lessons_model(received_prompts: list[str]) -> FunctionModel: + def respond(messages, _agent_info): + received_prompts.append(messages[-1].parts[-1].content) + return ModelResponse(parts=[TextPart(content=MOCK_ITERATIONS_LESSONS_RESPONSE)]) + + return FunctionModel(respond) + + +def _create_step(config, step_model, iterations_lessons_model) -> DataRepresentationStep: + return DataRepresentationStep(config, step_model, iterations_lessons_model, Mock(), []) + + +def test_iterations_lessons_from_successful_step_reach_next_step_prompt(initialized_run_config, prepared_iteration): + received_prompts: list[str] = [] + step = _create_step( + initialized_run_config, + create_model_calling_output_tool( + DataRepresentationOutput(representation="One-hot encoded nucleotide sequences.") + ), + _iterations_lessons_model(received_prompts), + ) + + asyncio.run(step.run()) + + assert len(received_prompts) == 1 + assert "One-hot encoded nucleotide sequences." in received_prompts[0] + next_step_prompt = _create_step(initialized_run_config, Mock(), Mock()).build_user_prompt() + assert f"\n{MOCK_ITERATIONS_LESSONS_RESPONSE}\n" in next_step_prompt + + +def test_iterations_lessons_from_failed_step_reach_next_step_prompt(initialized_run_config, prepared_iteration): + received_prompts: list[str] = [] + step = _create_step(initialized_run_config, _failing_model(), _iterations_lessons_model(received_prompts)) + + with pytest.raises(IterationRunFailed): + asyncio.run(step.run()) + + assert len(received_prompts) == 1 + assert "provider returned finish_reason='error'" in received_prompts[0] + next_step_prompt = _create_step(initialized_run_config, Mock(), Mock()).build_user_prompt() + assert f"\n{MOCK_ITERATIONS_LESSONS_RESPONSE}\n" in next_step_prompt