From 181a500c5d45d41a85e62461f4ee42b5442aef23 Mon Sep 17 00:00:00 2001
From: Albert Villanova del Moral
 <8515462+albertvillanova@users.noreply.github.com>
Date: Thu, 30 Jan 2025 21:24:39 +0100
Subject: [PATCH] Fix and test MemoryStep (#432)

* Test MemoryStep
* Remove unused MemoryStep.raw attribute
---
 src/smolagents/memory.py |  3 +--
 tests/test_memory.py     | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/smolagents/memory.py b/src/smolagents/memory.py
index 5b38e4893..80c8c7d5c 100644
--- a/src/smolagents/memory.py
+++ b/src/smolagents/memory.py
@@ -37,9 +37,8 @@ def dict(self):
         }
 
 
+@dataclass
 class MemoryStep:
-    raw: Any  # This is a placeholder for the raw data that the agent logs
-
     def dict(self):
         return asdict(self)
 
diff --git a/tests/test_memory.py b/tests/test_memory.py
index b2f2ffce7..0362272fa 100644
--- a/tests/test_memory.py
+++ b/tests/test_memory.py
@@ -1,7 +1,10 @@
+import pytest
+
 from smolagents.memory import (
     ActionStep,
     AgentMemory,
     ChatMessage,
+    MemoryStep,
     Message,
     MessageRole,
     PlanningStep,
@@ -18,6 +21,21 @@ def test_initialization(self):
         assert memory.steps == []
 
 
+class TestMemoryStep:
+    def test_initialization(self):
+        step = MemoryStep()
+        assert isinstance(step, MemoryStep)
+
+    def test_dict(self):
+        step = MemoryStep()
+        assert step.dict() == {}
+
+    def test_to_messages(self):
+        step = MemoryStep()
+        with pytest.raises(NotImplementedError):
+            step.to_messages()
+
+
 def test_action_step_to_messages():
     action_step = ActionStep(
         model_input_messages=[Message(role=MessageRole.USER, content="Hello")],