Fix and test MemoryStep (#432)

* Test MemoryStep
* Remove unused MemoryStep.raw attribute
This commit is contained in:
Albert Villanova del Moral 2025-01-30 21:24:39 +01:00 committed by GitHub
parent 42d97716fe
commit 181a500c5d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 19 additions and 2 deletions

View File

@ -37,9 +37,8 @@ class ToolCall:
}
@dataclass
class MemoryStep:
raw: Any # This is a placeholder for the raw data that the agent logs
def dict(self):
return asdict(self)

View File

@ -1,7 +1,10 @@
import pytest
from smolagents.memory import (
ActionStep,
AgentMemory,
ChatMessage,
MemoryStep,
Message,
MessageRole,
PlanningStep,
@ -18,6 +21,21 @@ class TestAgentMemory:
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")],