3131_TEST_LOCATION = "us-central1"
3232_TEST_PROJECT = "test-project"
3333_TEST_MODEL = "gemini-1.0-pro"
34+ _TEST_RUNNABLE_NAME = "test-runnable"
3435_TEST_SYSTEM_INSTRUCTION = "You are a helpful bot."
3536
3637
@@ -127,8 +128,11 @@ def teardown_method(self):
127128 initializer .global_pool .shutdown (wait = True )
128129
129130 def test_initialization (self ):
130- agent = reasoning_engines .AG2Agent (model = _TEST_MODEL )
131+ agent = reasoning_engines .AG2Agent (
132+ model = _TEST_MODEL , runnable_name = _TEST_RUNNABLE_NAME
133+ )
131134 assert agent ._model_name == _TEST_MODEL
135+ assert agent ._runnable_name == _TEST_RUNNABLE_NAME
132136 assert agent ._project == _TEST_PROJECT
133137 assert agent ._location == _TEST_LOCATION
134138 assert agent ._runnable is None
@@ -140,6 +144,7 @@ def test_initialization_with_tools(self, autogen_tools_mock):
140144 ]
141145 agent = reasoning_engines .AG2Agent (
142146 model = _TEST_MODEL ,
147+ runnable_name = _TEST_RUNNABLE_NAME ,
143148 system_instruction = _TEST_SYSTEM_INSTRUCTION ,
144149 tools = tools ,
145150 runnable_builder = lambda ** kwargs : kwargs ,
@@ -154,6 +159,7 @@ def test_initialization_with_tools(self, autogen_tools_mock):
154159 def test_set_up (self ):
155160 agent = reasoning_engines .AG2Agent (
156161 model = _TEST_MODEL ,
162+ runnable_name = _TEST_RUNNABLE_NAME ,
157163 runnable_builder = lambda ** kwargs : kwargs ,
158164 )
159165 assert agent ._runnable is None
@@ -163,6 +169,7 @@ def test_set_up(self):
163169 def test_clone (self ):
164170 agent = reasoning_engines .AG2Agent (
165171 model = _TEST_MODEL ,
172+ runnable_name = _TEST_RUNNABLE_NAME ,
166173 runnable_builder = lambda ** kwargs : kwargs ,
167174 )
168175 agent .set_up ()
@@ -176,6 +183,7 @@ def test_clone(self):
176183 def test_query (self , dataclasses_asdict_mock ):
177184 agent = reasoning_engines .AG2Agent (
178185 model = _TEST_MODEL ,
186+ runnable_name = _TEST_RUNNABLE_NAME ,
179187 )
180188 agent ._runnable = mock .Mock ()
181189 mocks = mock .Mock ()
@@ -202,6 +210,7 @@ def test_enable_tracing(
202210 ):
203211 agent = reasoning_engines .AG2Agent (
204212 model = _TEST_MODEL ,
213+ runnable_name = _TEST_RUNNABLE_NAME ,
205214 enable_tracing = True ,
206215 )
207216 assert agent ._enable_tracing is True
@@ -220,5 +229,6 @@ def test_raise_untyped_input_args(self, vertexai_init_mock):
220229 with pytest .raises (TypeError , match = r"has untyped input_arg" ):
221230 reasoning_engines .AG2Agent (
222231 model = _TEST_MODEL ,
232+ runnable_name = _TEST_RUNNABLE_NAME ,
223233 tools = [_return_input_no_typing ],
224234 )
0 commit comments