@@ -142,7 +142,7 @@ def memory_config(self):
142142 return MemoryConfig (memory_space_id = "test-space-id" )
143143
144144 def test_graph_has_memory_recall_node (
145- self , mock_model : MagicMock , messages : list , memory_config : MemoryConfig
145+ self , mock_model : MagicMock , messages : list [ SystemMessage ] , memory_config : MemoryConfig
146146 ) -> None :
147147 result : StateGraph [Any ] = create_agent (
148148 mock_model , [], messages , memory = memory_config
@@ -151,7 +151,7 @@ def test_graph_has_memory_recall_node(
151151 assert AgentGraphNode .MEMORY_RECALL in graph .nodes
152152
153153 def test_graph_edges_with_memory (
154- self , mock_model : MagicMock , messages : list , memory_config : MemoryConfig
154+ self , mock_model : MagicMock , messages : list [ SystemMessage ] , memory_config : MemoryConfig
155155 ) -> None :
156156 result : StateGraph [Any ] = create_agent (
157157 mock_model , [], messages , memory = memory_config
@@ -163,7 +163,7 @@ def test_graph_edges_with_memory(
163163 assert Edge ("__start__" , AgentGraphNode .INIT ) not in graph .edges
164164
165165 def test_graph_without_memory_has_no_recall_node (
166- self , mock_model : MagicMock , messages : list
166+ self , mock_model : MagicMock , messages : list [ SystemMessage ]
167167 ) -> None :
168168 result : StateGraph [Any ] = create_agent (mock_model , [], messages )
169169 graph = result .compile ().get_graph ()
0 commit comments