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 04acc9020487..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 @@ -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, diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py index 4565dd28d1d3..219443a21609 100644 --- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py +++ b/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": {