feat: support multiple workbenches in assistant agent

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
+    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 | None, optional): The workbench 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`.
@@ -651,6 +651,7 @@ async def run_reasoning_agent() -> None:
 
     """
 
+    component_version = 2
     component_config_schema = AssistantAgentConfig
     component_provider_override = "autogen_agentchat.agents.AssistantAgent"
 
@@ -660,7 +661,7 @@ def __init__(
         model_client: ChatCompletionClient,
         *,
         tools: List[BaseTool[Any, Any] | Callable[..., Any] | Callable[..., Awaitable[Any]]] | None = None,
-        workbench: 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,9 +749,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, Sequence):
+                self._workbench = workbench
+            else:
+                self._workbench = [workbench]
         else:
-            self._workbench = StaticWorkbench(self._tools)
+            self._workbench = [StaticWorkbench(self._tools)]
 
         if model_context is not None:
             self._model_context = model_context
@@ -942,7 +946,7 @@ async def _call_llm(
         model_client_stream: bool,
         system_messages: List[SystemMessage],
         model_context: ChatCompletionContext,
-        workbench: Workbench,
+        workbench: Sequence[Workbench],
         handoff_tools: List[BaseTool[Any, Any]],
         agent_name: str,
         cancellation_token: CancellationToken,
@@ -954,7 +958,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 = (await workbench.list_tools()) + handoff_tools
+        tools = [tool for wb in workbench for tool in await wb.list_tools()] + handoff_tools
 
         if model_client_stream:
             model_result: Optional[CreateResult] = None
@@ -991,7 +995,7 @@ async def _process_model_result(
         agent_name: str,
         system_messages: List[SystemMessage],
         model_context: ChatCompletionContext,
-        workbench: Workbench,
+        workbench: Sequence[Workbench],
         handoff_tools: List[BaseTool[Any, Any]],
         handoffs: Dict[str, HandoffBase],
         model_client: ChatCompletionClient,
@@ -1292,7 +1296,7 @@ def default_tool_call_summary_formatter(call: FunctionCall, result: FunctionExec
     @staticmethod
     async def _execute_tool_call(
         tool_call: FunctionCall,
-        workbench: Workbench,
+        workbench: Sequence[Workbench],
         handoff_tools: List[BaseTool[Any, Any]],
         agent_name: str,
         cancellation_token: CancellationToken,
@@ -1330,17 +1334,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 workbench:
+            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,
             ),
         )
@@ -1375,7 +1392,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=[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,
@@ -1407,7 +1424,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.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,

python/packages/autogen-agentchat/tests/test_assistant_agent.py

@@ -1483,7 +1483,9 @@ 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
+    assert agent.component_version == deserialize.component_version
 
 
 @pytest.mark.asyncio
@@ -1505,7 +1507,41 @@ 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
+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)  # type: ignore
+    assert len(deserialized_agent._workbench) == len(workbenches)  # type: ignore
+
+    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()  # type: ignore
 
 
 @pytest.mark.asyncio
@@ -1515,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": {