From 20392f52a6c08df9966051aec775a218e654afca Mon Sep 17 00:00:00 2001 From: Griffin Bassman Date: Wed, 14 May 2025 09:36:23 -0400 Subject: [PATCH 1/8] feat: support multiple workbenches in assistant agent --- .../agents/_assistant_agent.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py index c189021a2e2c..9f3652897344 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py @@ -642,7 +642,7 @@ def __init__( model_client: ChatCompletionClient, *, tools: List[BaseTool[Any, Any] | Callable[..., Any] | Callable[..., Awaitable[Any]]] | None = None, - workbench: Workbench | None = None, + workbench: Workbench | List[Workbench] | None = None, handoffs: List[HandoffBase | str] | None = None, model_context: ChatCompletionContext | None = None, description: str = "An agent that provides assistance with ability to use tools.", @@ -729,9 +729,12 @@ def __init__( if workbench is not None: if self._tools: raise ValueError("Tools cannot be used with a workbench.") - self._workbench = workbench + if isinstance(workbench, list): + self._workbenches = workbench + else: + self._workbenches = [workbench] else: - self._workbench = StaticWorkbench(self._tools) + self._workbenches = [StaticWorkbench(self._tools)] if model_context is not None: self._model_context = model_context @@ -796,7 +799,7 @@ async def on_messages_stream( model_context = self._model_context memory = self._memory system_messages = self._system_messages - workbench = self._workbench + workbenches = self._workbenches handoff_tools = self._handoff_tools handoffs = self._handoffs model_client = self._model_client @@ -829,7 +832,7 @@ async def on_messages_stream( model_client_stream=model_client_stream, system_messages=system_messages, model_context=model_context, - workbench=workbench, + workbench=workbenches[0], handoff_tools=handoff_tools, agent_name=agent_name, cancellation_token=cancellation_token, @@ -866,7 +869,7 @@ async def on_messages_stream( agent_name=agent_name, system_messages=system_messages, model_context=model_context, - workbench=workbench, + workbench=workbenches[0], handoff_tools=handoff_tools, handoffs=handoffs, model_client=model_client, @@ -1343,7 +1346,7 @@ def _to_config(self) -> AssistantAgentConfig: name=self.name, model_client=self._model_client.dump_component(), tools=None, # versionchanged:: v0.5.5 Now tools are not serialized, Cause they are part of the workbench. - workbench=self._workbench.dump_component() if self._workbench else None, + workbench=self._workbench.dump_component() if self._workbenches[0] else None, handoffs=list(self._handoffs.values()) if self._handoffs else None, model_context=self._model_context.dump_component(), memory=[memory.dump_component() for memory in self._memory] if self._memory else None, From 16471382684bdbeff3ece7035abac96188da72cc Mon Sep 17 00:00:00 2001 From: Griffin Bassman Date: Wed, 14 May 2025 10:10:52 -0400 Subject: [PATCH 2/8] only call workbench with needed tool --- .../agents/_assistant_agent.py | 52 ++++++++++++------- 1 file changed, 34 insertions(+), 18 deletions(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py index 9f3652897344..feb0d1ab2f16 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py @@ -66,7 +66,7 @@ class AssistantAgentConfig(BaseModel): name: str model_client: ComponentModel tools: List[ComponentModel] | None = None - workbench: ComponentModel | None = None + workbenches: List[ComponentModel] | None = None handoffs: List[HandoffBase | str] | None = None model_context: ComponentModel | None = None memory: List[ComponentModel] | None = None @@ -179,7 +179,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): name (str): The name of the agent. model_client (ChatCompletionClient): The model client to use for inference. tools (List[BaseTool[Any, Any] | Callable[..., Any] | Callable[..., Awaitable[Any]]] | None, optional): The tools to register with the agent. - workbench (Workbench | None, optional): The workbench to use for the agent. + workbench (Workbench | List[Workbench] | None, optional): The workbenches to use for the agent. Tools cannot be used when workbench is set and vice versa. handoffs (List[HandoffBase | str] | None, optional): The handoff configurations for the agent, allowing it to transfer to other agents by responding with a :class:`HandoffMessage`. @@ -832,7 +832,7 @@ async def on_messages_stream( model_client_stream=model_client_stream, system_messages=system_messages, model_context=model_context, - workbench=workbenches[0], + workbench=workbenches, handoff_tools=handoff_tools, agent_name=agent_name, cancellation_token=cancellation_token, @@ -869,7 +869,7 @@ async def on_messages_stream( agent_name=agent_name, system_messages=system_messages, model_context=model_context, - workbench=workbenches[0], + workbenches=workbenches, handoff_tools=handoff_tools, handoffs=handoffs, model_client=model_client, @@ -923,7 +923,7 @@ async def _call_llm( model_client_stream: bool, system_messages: List[SystemMessage], model_context: ChatCompletionContext, - workbench: Workbench, + workbenches: List[Workbench], handoff_tools: List[BaseTool[Any, Any]], agent_name: str, cancellation_token: CancellationToken, @@ -935,7 +935,10 @@ async def _call_llm( all_messages = await model_context.get_messages() llm_messages = cls._get_compatible_context(model_client=model_client, messages=system_messages + all_messages) - tools = (await workbench.list_tools()) + handoff_tools + tools = [] + for wb in workbenches: + tools.extend(await wb.list_tools()) + tools += handoff_tools if model_client_stream: model_result: Optional[CreateResult] = None @@ -972,7 +975,7 @@ async def _process_model_result( agent_name: str, system_messages: List[SystemMessage], model_context: ChatCompletionContext, - workbench: Workbench, + workbenches: List[Workbench], handoff_tools: List[BaseTool[Any, Any]], handoffs: Dict[str, HandoffBase], model_client: ChatCompletionClient, @@ -1031,7 +1034,7 @@ async def _process_model_result( *[ cls._execute_tool_call( tool_call=call, - workbench=workbench, + workbenches=workbenches, handoff_tools=handoff_tools, agent_name=agent_name, cancellation_token=cancellation_token, @@ -1263,7 +1266,7 @@ def _summarize_tool_use( @staticmethod async def _execute_tool_call( tool_call: FunctionCall, - workbench: Workbench, + workbenches: List[Workbench], handoff_tools: List[BaseTool[Any, Any]], agent_name: str, cancellation_token: CancellationToken, @@ -1301,17 +1304,30 @@ async def _execute_tool_call( ) # Handle normal tool call using workbench. - result = await workbench.call_tool( - name=tool_call.name, - arguments=arguments, - cancellation_token=cancellation_token, - ) + for wb in workbenches: + tools = await wb.list_tools() + if any(t["name"] == tool_call.name for t in tools): + result = await wb.call_tool( + name=tool_call.name, + arguments=arguments, + cancellation_token=cancellation_token, + ) + return ( + tool_call, + FunctionExecutionResult( + content=result.to_text(), + call_id=tool_call.id, + is_error=result.is_error, + name=tool_call.name, + ), + ) + return ( tool_call, FunctionExecutionResult( - content=result.to_text(), + content=f"Error: tool '{tool_call.name}' not found in any workbench", call_id=tool_call.id, - is_error=result.is_error, + is_error=True, name=tool_call.name, ), ) @@ -1346,7 +1362,7 @@ def _to_config(self) -> AssistantAgentConfig: name=self.name, model_client=self._model_client.dump_component(), tools=None, # versionchanged:: v0.5.5 Now tools are not serialized, Cause they are part of the workbench. - workbench=self._workbench.dump_component() if self._workbenches[0] else None, + workbenches=[wb.dump_component() for wb in self._workbenches] if self._workbenches else None, handoffs=list(self._handoffs.values()) if self._handoffs else None, model_context=self._model_context.dump_component(), memory=[memory.dump_component() for memory in self._memory] if self._memory else None, @@ -1378,7 +1394,7 @@ def _from_config(cls, config: AssistantAgentConfig) -> Self: return cls( name=config.name, model_client=ChatCompletionClient.load_component(config.model_client), - workbench=Workbench.load_component(config.workbench) if config.workbench else None, + workbench=[Workbench.load_component(wb) for wb in config.workbenches] if config.workbenches else None, handoffs=config.handoffs, model_context=ChatCompletionContext.load_component(config.model_context) if config.model_context else None, tools=[BaseTool.load_component(tool) for tool in config.tools] if config.tools else None, From 0f88e4b862d79e0845be1055d56d39a776f2d267 Mon Sep 17 00:00:00 2001 From: Griffin Bassman Date: Thu, 22 May 2025 09:37:42 -0400 Subject: [PATCH 3/8] fix workbenches in _call_llm --- .../src/autogen_agentchat/agents/_assistant_agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py index 7e86db46e249..cd201ebbf50d 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py @@ -853,7 +853,7 @@ async def on_messages_stream( model_client_stream=model_client_stream, system_messages=system_messages, model_context=model_context, - workbench=workbenches, + workbenches=workbenches, handoff_tools=handoff_tools, agent_name=agent_name, cancellation_token=cancellation_token, From 2af7d1211ec8d40dc24dfdfdca7f85b2833d0f26 Mon Sep 17 00:00:00 2001 From: Griffin Bassman Date: Thu, 22 May 2025 09:53:43 -0400 Subject: [PATCH 4/8] fix naming and test --- .../agents/_assistant_agent.py | 32 +++++++++---------- .../tests/test_assistant_agent.py | 31 ++++++++++++++++++ 2 files changed, 47 insertions(+), 16 deletions(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py index cd201ebbf50d..417b9b1437c9 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py @@ -66,7 +66,7 @@ class AssistantAgentConfig(BaseModel): name: str model_client: ComponentModel tools: List[ComponentModel] | None = None - workbenches: List[ComponentModel] | None = None + workbench: List[ComponentModel] | None = None handoffs: List[HandoffBase | str] | None = None model_context: ComponentModel | None = None memory: List[ComponentModel] | None = None @@ -188,7 +188,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): name (str): The name of the agent. model_client (ChatCompletionClient): The model client to use for inference. tools (List[BaseTool[Any, Any] | Callable[..., Any] | Callable[..., Awaitable[Any]]] | None, optional): The tools to register with the agent. - workbench (Workbench | List[Workbench] | None, optional): The workbenches to use for the agent. + workbench (Workbench | List[Workbench] | None, optional): The workbench or list of workbenches to use for the agent. Tools cannot be used when workbench is set and vice versa. handoffs (List[HandoffBase | str] | None, optional): The handoff configurations for the agent, allowing it to transfer to other agents by responding with a :class:`HandoffMessage`. @@ -749,11 +749,11 @@ def __init__( if self._tools: raise ValueError("Tools cannot be used with a workbench.") if isinstance(workbench, list): - self._workbenches = workbench + self._workbench = workbench else: - self._workbenches = [workbench] + self._workbench = [workbench] else: - self._workbenches = [StaticWorkbench(self._tools)] + self._workbench = [StaticWorkbench(self._tools)] if model_context is not None: self._model_context = model_context @@ -819,7 +819,7 @@ async def on_messages_stream( model_context = self._model_context memory = self._memory system_messages = self._system_messages - workbenches = self._workbenches + workbench = self._workbench handoff_tools = self._handoff_tools handoffs = self._handoffs model_client = self._model_client @@ -853,7 +853,7 @@ async def on_messages_stream( model_client_stream=model_client_stream, system_messages=system_messages, model_context=model_context, - workbenches=workbenches, + workbench=workbench, handoff_tools=handoff_tools, agent_name=agent_name, cancellation_token=cancellation_token, @@ -890,7 +890,7 @@ async def on_messages_stream( agent_name=agent_name, system_messages=system_messages, model_context=model_context, - workbenches=workbenches, + workbench=workbench, handoff_tools=handoff_tools, handoffs=handoffs, model_client=model_client, @@ -945,7 +945,7 @@ async def _call_llm( model_client_stream: bool, system_messages: List[SystemMessage], model_context: ChatCompletionContext, - workbenches: List[Workbench], + workbench: List[Workbench], handoff_tools: List[BaseTool[Any, Any]], agent_name: str, cancellation_token: CancellationToken, @@ -958,7 +958,7 @@ async def _call_llm( llm_messages = cls._get_compatible_context(model_client=model_client, messages=system_messages + all_messages) tools = [] - for wb in workbenches: + for wb in workbench: tools.extend(await wb.list_tools()) tools += handoff_tools @@ -997,7 +997,7 @@ async def _process_model_result( agent_name: str, system_messages: List[SystemMessage], model_context: ChatCompletionContext, - workbenches: List[Workbench], + workbench: List[Workbench], handoff_tools: List[BaseTool[Any, Any]], handoffs: Dict[str, HandoffBase], model_client: ChatCompletionClient, @@ -1057,7 +1057,7 @@ async def _process_model_result( *[ cls._execute_tool_call( tool_call=call, - workbenches=workbenches, + workbench=workbench, handoff_tools=handoff_tools, agent_name=agent_name, cancellation_token=cancellation_token, @@ -1298,7 +1298,7 @@ def default_tool_call_summary_formatter(call: FunctionCall, result: FunctionExec @staticmethod async def _execute_tool_call( tool_call: FunctionCall, - workbenches: List[Workbench], + workbench: List[Workbench], handoff_tools: List[BaseTool[Any, Any]], agent_name: str, cancellation_token: CancellationToken, @@ -1336,7 +1336,7 @@ async def _execute_tool_call( ) # Handle normal tool call using workbench. - for wb in workbenches: + for wb in workbench: tools = await wb.list_tools() if any(t["name"] == tool_call.name for t in tools): result = await wb.call_tool( @@ -1394,7 +1394,7 @@ def _to_config(self) -> AssistantAgentConfig: name=self.name, model_client=self._model_client.dump_component(), tools=None, # versionchanged:: v0.5.5 Now tools are not serialized, Cause they are part of the workbench. - workbenches=[wb.dump_component() for wb in self._workbenches] if self._workbenches else None, + workbench=[wb.dump_component() for wb in self._workbench] if self._workbench else None, handoffs=list(self._handoffs.values()) if self._handoffs else None, model_context=self._model_context.dump_component(), memory=[memory.dump_component() for memory in self._memory] if self._memory else None, @@ -1426,7 +1426,7 @@ def _from_config(cls, config: AssistantAgentConfig) -> Self: return cls( name=config.name, model_client=ChatCompletionClient.load_component(config.model_client), - workbench=[Workbench.load_component(wb) for wb in config.workbenches] if config.workbenches else None, + workbench=[Workbench.load_component(wb) for wb in config.workbench] if config.workbench else None, handoffs=config.handoffs, model_context=ChatCompletionContext.load_component(config.model_context) if config.model_context else None, tools=[BaseTool.load_component(tool) for tool in config.tools] if config.tools else None, diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py index 988707d9549b..8948332aa4c3 100644 --- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py +++ b/python/packages/autogen-agentchat/tests/test_assistant_agent.py @@ -1508,6 +1508,37 @@ async def test_workbenchs_serialize_and_deserialize() -> None: assert deserialize._workbench._to_config() == agent._workbench._to_config() # type: ignore +@pytest.mark.asyncio +async def test_multiple_workbenchs_serialize_and_deserialize() -> None: + workbenches: List[McpWorkbench] = [ + McpWorkbench(server_params=SseServerParams(url="http://test-url-1")), + McpWorkbench(server_params=SseServerParams(url="http://test-url-2")), + ] + + client = OpenAIChatCompletionClient( + model="gpt-4o", + api_key="API_KEY", + ) + + agent = AssistantAgent( + name="test_multi", + model_client=client, + workbench=workbenches, + ) + + serialize = agent.dump_component() + deserialized_agent: AssistantAgent = AssistantAgent.load_component(serialize) + + assert deserialized_agent.name == agent.name + assert isinstance(deserialized_agent._workbench, list) + assert len(deserialized_agent._workbench) == len(workbenches) + + for original, restored in zip(agent._workbench, deserialized_agent._workbench, strict=True): + assert isinstance(original, McpWorkbench) + assert isinstance(restored, McpWorkbench) + assert original._to_config() == restored._to_config() + + @pytest.mark.asyncio async def test_tools_deserialize_aware() -> None: dump = """ From 0bbbd44121a965f95e75c20295842844cba6c1ce Mon Sep 17 00:00:00 2001 From: Griffin Bassman Date: Thu, 22 May 2025 10:11:33 -0400 Subject: [PATCH 5/8] fix mypy --- .../src/autogen_agentchat/agents/_assistant_agent.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py index 417b9b1437c9..5383ae8f7796 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py @@ -957,10 +957,7 @@ async def _call_llm( all_messages = await model_context.get_messages() llm_messages = cls._get_compatible_context(model_client=model_client, messages=system_messages + all_messages) - tools = [] - for wb in workbench: - tools.extend(await wb.list_tools()) - tools += handoff_tools + tools = [tool for wb in workbench for tool in await wb.list_tools()] + handoff_tools # type: ignore[arg-type] if model_client_stream: model_result: Optional[CreateResult] = None From da47797967c22af0635f2fc26ae87935212b5856 Mon Sep 17 00:00:00 2001 From: Griffin Bassman Date: Thu, 22 May 2025 10:35:08 -0400 Subject: [PATCH 6/8] Sequence to fix pyright issues --- .../autogen_agentchat/agents/_assistant_agent.py | 14 +++++++------- .../tests/test_assistant_agent.py | 8 ++++---- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py index 5383ae8f7796..9f8e8d710cdf 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py @@ -188,7 +188,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]): name (str): The name of the agent. model_client (ChatCompletionClient): The model client to use for inference. tools (List[BaseTool[Any, Any] | Callable[..., Any] | Callable[..., Awaitable[Any]]] | None, optional): The tools to register with the agent. - workbench (Workbench | List[Workbench] | None, optional): The workbench or list of workbenches to use for the agent. + workbench (Workbench | Sequence[Workbench] | None, optional): The workbench or list of workbenches to use for the agent. Tools cannot be used when workbench is set and vice versa. handoffs (List[HandoffBase | str] | None, optional): The handoff configurations for the agent, allowing it to transfer to other agents by responding with a :class:`HandoffMessage`. @@ -660,7 +660,7 @@ def __init__( model_client: ChatCompletionClient, *, tools: List[BaseTool[Any, Any] | Callable[..., Any] | Callable[..., Awaitable[Any]]] | None = None, - workbench: Workbench | List[Workbench] | None = None, + workbench: Workbench | Sequence[Workbench] | None = None, handoffs: List[HandoffBase | str] | None = None, model_context: ChatCompletionContext | None = None, description: str = "An agent that provides assistance with ability to use tools.", @@ -748,7 +748,7 @@ def __init__( if workbench is not None: if self._tools: raise ValueError("Tools cannot be used with a workbench.") - if isinstance(workbench, list): + if isinstance(workbench, Sequence): self._workbench = workbench else: self._workbench = [workbench] @@ -945,7 +945,7 @@ async def _call_llm( model_client_stream: bool, system_messages: List[SystemMessage], model_context: ChatCompletionContext, - workbench: List[Workbench], + workbench: Sequence[Workbench], handoff_tools: List[BaseTool[Any, Any]], agent_name: str, cancellation_token: CancellationToken, @@ -957,7 +957,7 @@ async def _call_llm( all_messages = await model_context.get_messages() llm_messages = cls._get_compatible_context(model_client=model_client, messages=system_messages + all_messages) - tools = [tool for wb in workbench for tool in await wb.list_tools()] + handoff_tools # type: ignore[arg-type] + tools = [tool for wb in workbench for tool in await wb.list_tools()] + handoff_tools if model_client_stream: model_result: Optional[CreateResult] = None @@ -994,7 +994,7 @@ async def _process_model_result( agent_name: str, system_messages: List[SystemMessage], model_context: ChatCompletionContext, - workbench: List[Workbench], + workbench: Sequence[Workbench], handoff_tools: List[BaseTool[Any, Any]], handoffs: Dict[str, HandoffBase], model_client: ChatCompletionClient, @@ -1295,7 +1295,7 @@ def default_tool_call_summary_formatter(call: FunctionCall, result: FunctionExec @staticmethod async def _execute_tool_call( tool_call: FunctionCall, - workbench: List[Workbench], + workbench: Sequence[Workbench], handoff_tools: List[BaseTool[Any, Any]], agent_name: str, cancellation_token: CancellationToken, diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py index 8948332aa4c3..50c5333dde64 100644 --- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py +++ b/python/packages/autogen-agentchat/tests/test_assistant_agent.py @@ -1530,13 +1530,13 @@ async def test_multiple_workbenchs_serialize_and_deserialize() -> None: deserialized_agent: AssistantAgent = AssistantAgent.load_component(serialize) assert deserialized_agent.name == agent.name - assert isinstance(deserialized_agent._workbench, list) - assert len(deserialized_agent._workbench) == len(workbenches) + assert isinstance(deserialized_agent._workbench, list) # type: ignore + assert len(deserialized_agent._workbench) == len(workbenches) # type: ignore - for original, restored in zip(agent._workbench, deserialized_agent._workbench, strict=True): + for original, restored in zip(agent._workbench, deserialized_agent._workbench, strict=True): # type: ignore assert isinstance(original, McpWorkbench) assert isinstance(restored, McpWorkbench) - assert original._to_config() == restored._to_config() + assert original._to_config() == restored._to_config() # type: ignore @pytest.mark.asyncio From 5d1a7abb2048c9b12a648993f8cfb90883bb385a Mon Sep 17 00:00:00 2001 From: Griffin Bassman Date: Thu, 22 May 2025 10:55:25 -0400 Subject: [PATCH 7/8] fix tools --- .../autogen-agentchat/tests/test_assistant_agent.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py index 50c5333dde64..75ed26943656 100644 --- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py +++ b/python/packages/autogen-agentchat/tests/test_assistant_agent.py @@ -1483,7 +1483,8 @@ def test() -> str: deserialize = AssistantAgent.load_component(serialize) assert deserialize.name == agent.name - assert await deserialize._workbench.list_tools() == await agent._workbench.list_tools() # type: ignore + for original, restored in zip(agent._workbench, deserialize._workbench, strict=True): # type: ignore + assert await original.list_tools() == await restored.list_tools() # type: ignore @pytest.mark.asyncio @@ -1505,7 +1506,10 @@ async def test_workbenchs_serialize_and_deserialize() -> None: deserialize = AssistantAgent.load_component(serialize) assert deserialize.name == agent.name - assert deserialize._workbench._to_config() == agent._workbench._to_config() # type: ignore + for original, restored in zip(agent._workbench, deserialize._workbench, strict=True): # type: ignore + assert isinstance(original, McpWorkbench) + assert isinstance(restored, McpWorkbench) + assert original._to_config() == restored._to_config() # type: ignore @pytest.mark.asyncio From a2ef69961bfb2d665b9359caab95a62c8f7ff9bf Mon Sep 17 00:00:00 2001 From: Griffin Bassman Date: Wed, 28 May 2025 13:23:34 -0400 Subject: [PATCH 8/8] add assistantagent component version 2 --- .../src/autogen_agentchat/agents/_assistant_agent.py | 1 + .../packages/autogen-agentchat/tests/test_assistant_agent.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py index 9f8e8d710cdf..12995b8ca150 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py @@ -651,6 +651,7 @@ async def run_reasoning_agent() -> None: """ + component_version = 2 component_config_schema = AssistantAgentConfig component_provider_override = "autogen_agentchat.agents.AssistantAgent" diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py index f3a97b2c4cd2..219443a21609 100644 --- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py +++ b/python/packages/autogen-agentchat/tests/test_assistant_agent.py @@ -1485,6 +1485,7 @@ def test() -> str: assert deserialize.name == agent.name for original, restored in zip(agent._workbench, deserialize._workbench, strict=True): # type: ignore assert await original.list_tools() == await restored.list_tools() # type: ignore + assert agent.component_version == deserialize.component_version @pytest.mark.asyncio @@ -1550,7 +1551,7 @@ async def test_tools_deserialize_aware() -> None: "provider": "autogen_agentchat.agents.AssistantAgent", "component_type": "agent", "version": 1, - "component_version": 1, + "component_version": 2, "description": "An agent that provides assistance with tool use.", "label": "AssistantAgent", "config": {