Fix and test MemoryStep (#432)
* Test MemoryStep * Remove unused MemoryStep.raw attribute
This commit is contained in:
parent
42d97716fe
commit
181a500c5d
|
@ -37,9 +37,8 @@ class ToolCall:
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
class MemoryStep:
|
class MemoryStep:
|
||||||
raw: Any # This is a placeholder for the raw data that the agent logs
|
|
||||||
|
|
||||||
def dict(self):
|
def dict(self):
|
||||||
return asdict(self)
|
return asdict(self)
|
||||||
|
|
||||||
|
|
|
@ -1,7 +1,10 @@
|
||||||
|
import pytest
|
||||||
|
|
||||||
from smolagents.memory import (
|
from smolagents.memory import (
|
||||||
ActionStep,
|
ActionStep,
|
||||||
AgentMemory,
|
AgentMemory,
|
||||||
ChatMessage,
|
ChatMessage,
|
||||||
|
MemoryStep,
|
||||||
Message,
|
Message,
|
||||||
MessageRole,
|
MessageRole,
|
||||||
PlanningStep,
|
PlanningStep,
|
||||||
|
@ -18,6 +21,21 @@ class TestAgentMemory:
|
||||||
assert memory.steps == []
|
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():
|
def test_action_step_to_messages():
|
||||||
action_step = ActionStep(
|
action_step = ActionStep(
|
||||||
model_input_messages=[Message(role=MessageRole.USER, content="Hello")],
|
model_input_messages=[Message(role=MessageRole.USER, content="Hello")],
|
||||||
|
|
Loading…
Reference in New Issue