SakanaAI/ale_agent

"""Rolling search-trajectory summary for ALE-Agent ablations."""
from __future__ import annotations

from ...components.memory import Memory


class ALEAgentMemory(Memory):
    def __init__(self, max_entries: int = 12):
        self.max_entries = max(1, int(max_entries))
        self._entries: list[str] = []

    def read(self, spec: dict | None = None) -> str:
        top = max(1, int((spec or {}).get("top", self.max_entries)))
        return "\n".join(f"- {entry}" for entry in self._entries[-top:])

    def write(self, knowledge: str, **meta) -> None:
        value = str(knowledge or "").strip()
        if value:
            self._entries.append(value)
            self._entries = self._entries[-self.max_entries:]

    def state_dict(self) -> dict:
        return {"entries": list(self._entries)}

    def load_state_dict(self, state: dict) -> None:
        if isinstance(state, dict) and isinstance(state.get("entries"), list):
            self._entries = [str(entry) for entry in state["entries"]][-self.max_entries:]