Allow dict input for state

Author: Unshure

src/strands/agent/agent.py

@@ -224,7 +224,7 @@ def __init__(
         *,
         name: Optional[str] = None,
         description: Optional[str] = None,
-        state: Optional[AgentState] = None,
+        state: Optional[Union[AgentState, dict]] = None,
     ):
         """Initialize the Agent with the specified configuration.
 
@@ -261,7 +261,7 @@ def __init__(
                 Defaults to None.
             description: description of what the Agent does
                 Defaults to None.
-            state: stateful information for the agent
+            state: stateful information for the agent. Can be either an AgentState object, or a json serializable dict.
                 Defaults to an empty AgentState object.
 
         Raises:
@@ -325,7 +325,17 @@ def __init__(
         self.trace_span: Optional[trace.Span] = None
 
         # Initialize agent state management
-        self.state = state or AgentState()
+        if state is not None:
+            if isinstance(state, dict):
+                self.state = AgentState(state)
+            elif isinstance(state, AgentState):
+                print("HERE!")
+                print(type(state))
+                self.state = state
+            else:
+                raise ValueError("state must be an AgentState object or a dict")
+        else:
+            self.state = AgentState()
 
         self.tool_caller = Agent.ToolCaller(self)
         self.name = name

tests/strands/agent/test_agent.py

@@ -1314,3 +1314,15 @@ def test_event_loop_cycle_includes_parent_span(mock_get_tracer, mock_event_loop_
     kwargs = mock_event_loop_cycle.call_args[1]
     assert "event_loop_parent_span" in kwargs
     assert kwargs["event_loop_parent_span"] == mock_span
+
+
+def test_non_dict_throws_error():
+    with pytest.raises(ValueError, match="state must be an AgentState object or a dict"):
+        agent = Agent(state={"object", object()})
+        print(agent.state)
+
+
+def test_non_json_serializable_state_throws_error():
+    with pytest.raises(ValueError, match="Value is not JSON serializable"):
+        agent = Agent(state={"object": object()})
+        print(agent.state)

tests/strands/mocked_model_provider/test_agent_state_updates.py

@@ -8,6 +8,7 @@
 @tool
 def update_state(agent: Agent):
     agent.state.set("hello", "world")
+    agent.state.set("foo", "baz")
 
 
 def test_agent_state_update_from_tool():
@@ -20,10 +21,16 @@ def test_agent_state_update_from_tool():
     ]
     mocked_model_provider = MockedModelProvider(agent_messages)
 
-    agent = Agent(model=mocked_model_provider, tools=[update_state])
+    agent = Agent(
+        model=mocked_model_provider,
+        tools=[update_state],
+        state={"foo": "bar"},
+    )
 
     assert agent.state.get("hello") is None
+    assert agent.state.get("foo") == "bar"
 
     agent("Invoke Mocked!")
 
     assert agent.state.get("hello") == "world"
+    assert agent.state.get("foo") == "baz"