feat: cross-iteration memory

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.lessons import read_lessons, update_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,18 @@ def on_step_start(self) -> None:
             check=True,
         )
 
+    async def on_step_end(self, error: Exception | None = None) -> None:
+        await update_lessons(
+            config=self.config,
+            model=self.iteration_plan_model,
+            step_name=self.display_name,
+            messages=self.step_messages(error),
+            error=error,
+        )
+
+    def step_messages(self, error: Exception | None) -> list[ModelMessage]:
+        return []
+
     def on_step_success(self, output: Any) -> None:
         save_step_output(self.config, self.step_id, output)
         export_shared_environment_descriptor(self.config)
@@ -116,6 +136,8 @@ class AgenticStepOutput(BaseModel):
 
 class AgenticStep(RuntimeStep):
     output_type: type[AgenticStepOutput]
+    _history_length: int = 0
+    _run_messages: list[ModelMessage] = []
 
     def on_step_start(self) -> None:
         super().on_step_start()
@@ -185,6 +207,13 @@ 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."
+        lessons = read_lessons(self.config)
+        # re-check: the whole lessons file goes to every step, consider sending only the sections relevant to this step
+        if lessons:
+            prompt += (
+                "\nLessons established so far in this run. Apply the ones relevant to this step instead of re-discovering them.\n"
+                f"<iterations_lessons>\n{lessons}\n</iterations_lessons>"
+            )
         return prompt
 
     def build_deps(self, step_started_at: datetime) -> dict[str, Any]:
@@ -215,12 +244,17 @@ def get_message_history(self):
             )
         return history
 
+    def step_messages(self, error: Exception | None) -> list[ModelMessage]:
+        messages = error.context_messages if isinstance(error, IterationRunFailed) else self._run_messages
+        return messages[self._history_length:]
+
     async def _execute(self) -> AgenticStepOutput:
         step_started_at = datetime.now()
         message_history = self.get_message_history()
+        self._history_length = len(message_history or [])
         user_prompt = self.build_user_prompt()
         agent = self.create_agent()
-        _, output = await run_agent(
+        self._run_messages, output = await run_agent(
             agent=agent,
             user_prompt=user_prompt,
             max_steps=self.config.max_steps,

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.lessons import read_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()
+        lessons_info = self._build_lessons_info()
 
         return f"""
         <prompt_of_the_iteration_agents>
@@ -114,6 +116,7 @@ def step_prompt(self) -> str:
         </dataset_knowledge_from_dataset_description_md>
 
         {iteration_history_info}
+        {lessons_info}
         <your_instructions>
         {f"<exploration_guidance>{exploration_info}</exploration_guidance>" 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,18 @@ def step_prompt(self) -> str:
     def get_message_history(self):
         return None
 
+    def _build_lessons_info(self) -> str:
+        lessons = read_lessons(self.config)
+        # re-check: the planner gets the whole lessons file, consider sending only the sections relevant to planning
+        if not 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>
+        {lessons}
+        </iterations_lessons>
+        """
+
     def _build_iteration_history_info(self) -> str:
         sorted_iteration_to_records = {
             iteration: self._load_iteration_history_record(iteration)

src/agentomics/runtime/lessons.py

@@ -0,0 +1,137 @@
+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
+
+LESSONS_AGENT_MAX_REQUESTS = 5
+NO_NEW_LESSONS = "NO NEW LESSONS"
+MAX_PART_CHARS = 1500
+MAX_TRANSCRIPT_CHARS = 60000
+
+# re-check: parts not listed here (system prompt, user prompt, thinking) are left out of the transcript
+# re-check: 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,
+}
+
+LESSONS_SYSTEM_PROMPT = """
+You maintain iterations_lessons.md, the memory of mistakes made by an autonomous ML agent.
+The agent 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 the step that just ended. You return the complete new content of the file.
+
+A lesson is a mistake the agent made and what it should do instead:
+- an error or dead end and its root cause
+- the fix or workaround that eventually worked
+- a validator complaint and how it was satisfied
+- environment problems, such as a missing package, a timeout, or a wrong path
+- workflow problems, such as repeating the same tool call or not finishing the step
+A fact about the data counts only if it caused a mistake, for example "ids in labels.csv are zero-padded, so joining on raw ids produced no rows".
+Don't record findings about the data or the model, good practices that caused no mistake, plans, or generic ML advice.
+Never invent anything the transcript doesn't show.
+
+Write each lesson as an instruction for the next agent, in this shape:
+- <Step or "Any step">: <what to do or avoid>. <short reason>.
+For example:
+- Exploration: labels.csv has `id` and `numeric_label`, no `diagnosis`. Join data on `id`.
+- Any step: don't repeat `ls` or `pwd`. Once your files are written, call `final_result`.
+Keep the names the next agent needs, such as columns, packages, and functions.
+Leave out full error messages, absolute paths, script names, and the story of how the fix was found.
+
+Editing rules:
+- Keep every existing lesson unless the transcript shows it is wrong. In that case, correct or remove it.
+- If the file already states the same mistake in other words, merge the new detail into that bullet instead of adding a near-duplicate.
+- If the step made no new mistakes, return exactly: NO NEW LESSONS
+- Keep each lesson under about 25 words and the file under 60 bullets.
+- Output only the markdown content of the file, without code fences or commentary.
+"""
+
+
+def read_lessons(config: Config) -> str:
+    if not config.lessons_path.exists():
+        return ""
+    return config.lessons_path.read_text(encoding="utf-8").strip()
+
+
+async def update_lessons(
+    config: Config,
+    model: Model,
+    step_name: str,
+    messages: list[ModelMessage],
+    error: Exception | None = None,
+) -> None:
+    try:
+        _, lessons = await run_agent(
+            agent=_create_lessons_agent(config, model),
+            user_prompt=_build_user_prompt(config, step_name, messages, error),
+            max_steps=LESSONS_AGENT_MAX_REQUESTS,
+            message_history=None,
+            verbose=False,
+        )
+    except Exception:
+        # ? skip instead of failing the run
+        print(f"WARNING: {Config.LESSONS_FILENAME} was not updated after step {step_name}.")
+        return
+    if lessons.strip() != NO_NEW_LESSONS:
+        config.lessons_path.write_text(f"{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.LESSONS_FILENAME}:
+{read_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:
+    lines = []
+    for message in messages:
+        for part in message.parts:
+            extract_content = _PART_CONTENT.get(type(part))
+            if extract_content is None:
+                continue
+            content = extract_content(part)
+            if content.strip():
+                lines.append(f"[{type(part).__name__}] {content[:MAX_PART_CHARS]}")
+    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_lessons_agent(config: Config, model: Model) -> Agent[None, str]:
+    return Agent(
+        model=model,
+        system_prompt=LESSONS_SYSTEM_PROMPT,
+        model_settings=ModelSettings(temperature=config.temperature),
+        output_type=str,
+        retries=config.max_validation_retries,
+    )

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"
+    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 lessons_path(self) -> Path:
+        return self.shared_dir / self.LESSONS_FILENAME
+
     @property
     def config_path(self) -> Path:
         return self.shared_dir / self.CONFIG_FILENAME

tests/integration/test_step_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
+
+LESSONS = "## Representation\n- One-hot encoding fits the fixed sequence length."
+
+
+def _failing_model() -> FunctionModel:
+    def respond(_messages, _agent_info):
+        raise RuntimeError("provider returned finish_reason='error'")
+
+    return FunctionModel(respond)
+
+
+def _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=LESSONS)])
+
+    return FunctionModel(respond)
+
+
+def _create_step(config, step_model, lessons_model) -> DataRepresentationStep:
+    return DataRepresentationStep(config, step_model, lessons_model, Mock(), [])
+
+
+def test_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.")
+        ),
+        _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"<iterations_lessons>\n{LESSONS}\n</iterations_lessons>" in next_step_prompt
+
+
+def test_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(), _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"<iterations_lessons>\n{LESSONS}\n</iterations_lessons>" in next_step_prompt