fix: improve workflow add_agent arguments (#518)

* fix: improve add_agent Signed-off-by: va <va@us.ibm.com> * fix: review feedback Signed-off-by: va <va@us.ibm.com> --------- Signed-off-by: va <va@us.ibm.com>
i-am-bee · Mar 7, 2025 · ee58c5f · ee58c5f
1 parent ab9dcbe
commit ee58c5f
Show file tree

Hide file tree

Showing 7 changed files with 117 additions and 98 deletions.
diff --git a/python/README.md b/python/README.md
@@ -89,29 +89,23 @@ async def main() -> None:
     try:
         workflow = AgentWorkflow(name="Smart assistant")
         workflow.add_agent(
-            agent=AgentFactoryInput(
-                name="WeatherForecaster",
-                instructions="You are a weather assistant. Respond only if you can provide a useful answer.",
-                tools=[OpenMeteoTool()],
-                llm=llm,
-                execution=AgentExecutionConfig(max_iterations=3),
-            )
+            name="WeatherForecaster",
+            instructions="You are a weather assistant. Respond only if you can provide a useful answer.",
+            tools=[OpenMeteoTool()],
+            llm=llm,
+            execution=AgentExecutionConfig(max_iterations=3),
         )
         workflow.add_agent(
-            agent=AgentFactoryInput(
-                name="Researcher",
-                instructions="You are a researcher assistant. Respond only if you can provide a useful answer.",
-                tools=[DuckDuckGoSearchTool()],
-                llm=llm,
-            )
+            name="Researcher",
+            instructions="You are a researcher assistant. Respond only if you can provide a useful answer.",
+            tools=[DuckDuckGoSearchTool()],
+            llm=llm,
         )
         workflow.add_agent(
-            agent=AgentFactoryInput(
-                name="Solver",
-                instructions="""Your task is to provide the most useful final answer based on the assistants'
+            name="Solver",
+            instructions="""Your task is to provide the most useful final answer based on the assistants'
 responses which all are relevant. Ignore those where assistant do not know.""",
-                llm=llm,
-            )
+            llm=llm,
         )
 
         prompt = "What is the weather in New York?"

diff --git a/python/beeai_framework/workflows/agent.py b/python/beeai_framework/workflows/agent.py
@@ -17,7 +17,7 @@
 import string
 from collections.abc import Awaitable, Callable
 from inspect import isfunction
-from typing import Self
+from typing import Any, Self
 
 from pydantic import BaseModel, ConfigDict, Field, InstanceOf
 
@@ -69,8 +69,22 @@ def del_agent(self, name: str) -> "AgentWorkflow":
 
     def add_agent(
         self,
-        agent: (BaseAgent | Callable[[ReadOnlyMemory], BaseAgent | asyncio.Future[BaseAgent]] | AgentFactoryInput),
+        agent: (
+            BaseAgent | Callable[[ReadOnlyMemory], BaseAgent | asyncio.Future[BaseAgent]] | AgentFactoryInput | None
+        ) = None,
+        /,
+        **kwargs: Any,
     ) -> "AgentWorkflow":
+        if not agent:
+            if not kwargs:
+                raise ValueError("An agent object or keyword arguments must be provided")
+            elif "agent" in kwargs:
+                agent = kwargs.get("agent")
+            else:
+                agent = AgentFactoryInput.model_validate(kwargs, strict=False, from_attributes=True)
+        elif kwargs:
+            raise ValueError("Agent object required or keyword arguments required but not both")
+
         if isinstance(agent, BaseAgent):
 
             async def factory(memory: ReadOnlyMemory) -> BaseAgent:

diff --git a/python/docs/agents.md b/python/docs/agents.md
@@ -452,37 +452,31 @@ from beeai_framework.errors import FrameworkError
 from beeai_framework.memory import UnconstrainedMemory
 from beeai_framework.tools.search.duckduckgo import DuckDuckGoSearchTool
 from beeai_framework.tools.weather.openmeteo import OpenMeteoTool
-from beeai_framework.workflows.agent import AgentFactoryInput, AgentWorkflow
+from beeai_framework.workflows.agent import AgentWorkflow
 
 
 async def main() -> None:
     llm = ChatModel.from_name("ollama:granite3.1-dense:8b")
 
     workflow = AgentWorkflow(name="Smart assistant")
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="WeatherForecaster",
-            instructions="You are a weather assistant.",
-            tools=[OpenMeteoTool()],
-            llm=llm,
-            execution=AgentExecutionConfig(max_iterations=3, total_max_retries=10, max_retries_per_step=3),
-        )
+        name="WeatherForecaster",
+        instructions="You are a weather assistant.",
+        tools=[OpenMeteoTool()],
+        llm=llm,
+        execution=AgentExecutionConfig(max_iterations=3, total_max_retries=10, max_retries_per_step=3),
     )
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="Researcher",
-            instructions="You are a researcher assistant.",
-            tools=[DuckDuckGoSearchTool()],
-            llm=llm,
-        )
+        name="Researcher",
+        instructions="You are a researcher assistant.",
+        tools=[DuckDuckGoSearchTool()],
+        llm=llm,
     )
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="Solver",
-            instructions="""Your task is to provide the most useful final answer based on the assistants'
+        name="Solver",
+        instructions="""Your task is to provide the most useful final answer based on the assistants'
 responses which all are relevant. Ignore those where assistant do not know.""",
-            llm=llm,
-        )
+        llm=llm,
     )
 
     prompt = "What is the weather in New York?"

diff --git a/python/docs/workflows.md b/python/docs/workflows.md
@@ -277,37 +277,31 @@ from beeai_framework.errors import FrameworkError
 from beeai_framework.memory import UnconstrainedMemory
 from beeai_framework.tools.search.duckduckgo import DuckDuckGoSearchTool
 from beeai_framework.tools.weather.openmeteo import OpenMeteoTool
-from beeai_framework.workflows.agent import AgentFactoryInput, AgentWorkflow
+from beeai_framework.workflows.agent import AgentWorkflow
 
 
 async def main() -> None:
     llm = ChatModel.from_name("ollama:granite3.1-dense:8b")
 
     workflow = AgentWorkflow(name="Smart assistant")
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="WeatherForecaster",
-            instructions="You are a weather assistant.",
-            tools=[OpenMeteoTool()],
-            llm=llm,
-            execution=AgentExecutionConfig(max_iterations=3, total_max_retries=10, max_retries_per_step=3),
-        )
+        name="WeatherForecaster",
+        instructions="You are a weather assistant.",
+        tools=[OpenMeteoTool()],
+        llm=llm,
+        execution=AgentExecutionConfig(max_iterations=3, total_max_retries=10, max_retries_per_step=3),
     )
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="Researcher",
-            instructions="You are a researcher assistant.",
-            tools=[DuckDuckGoSearchTool()],
-            llm=llm,
-        )
+        name="Researcher",
+        instructions="You are a researcher assistant.",
+        tools=[DuckDuckGoSearchTool()],
+        llm=llm,
     )
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="Solver",
-            instructions="""Your task is to provide the most useful final answer based on the assistants'
+        name="Solver",
+        instructions="""Your task is to provide the most useful final answer based on the assistants'
 responses which all are relevant. Ignore those where assistant do not know.""",
-            llm=llm,
-        )
+        llm=llm,
     )
 
     prompt = "What is the weather in New York?"

diff --git a/python/examples/workflows/multi_agents.py b/python/examples/workflows/multi_agents.py
@@ -9,37 +9,31 @@
 from beeai_framework.memory import UnconstrainedMemory
 from beeai_framework.tools.search.duckduckgo import DuckDuckGoSearchTool
 from beeai_framework.tools.weather.openmeteo import OpenMeteoTool
-from beeai_framework.workflows.agent import AgentFactoryInput, AgentWorkflow
+from beeai_framework.workflows.agent import AgentWorkflow
 
 
 async def main() -> None:
     llm = ChatModel.from_name("ollama:granite3.1-dense:8b")
 
     workflow = AgentWorkflow(name="Smart assistant")
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="WeatherForecaster",
-            instructions="You are a weather assistant.",
-            tools=[OpenMeteoTool()],
-            llm=llm,
-            execution=AgentExecutionConfig(max_iterations=3, total_max_retries=10, max_retries_per_step=3),
-        )
+        name="WeatherForecaster",
+        instructions="You are a weather assistant.",
+        tools=[OpenMeteoTool()],
+        llm=llm,
+        execution=AgentExecutionConfig(max_iterations=3, total_max_retries=10, max_retries_per_step=3),
     )
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="Researcher",
-            instructions="You are a researcher assistant.",
-            tools=[DuckDuckGoSearchTool()],
-            llm=llm,
-        )
+        name="Researcher",
+        instructions="You are a researcher assistant.",
+        tools=[DuckDuckGoSearchTool()],
+        llm=llm,
     )
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="Solver",
-            instructions="""Your task is to provide the most useful final answer based on the assistants'
+        name="Solver",
+        instructions="""Your task is to provide the most useful final answer based on the assistants'
 responses which all are relevant. Ignore those where assistant do not know.""",
-            llm=llm,
-        )
+        llm=llm,
     )
 
     prompt = "What is the weather in New York?"

diff --git a/python/examples/workflows/multi_agents_simple.py b/python/examples/workflows/multi_agents_simple.py
@@ -9,37 +9,31 @@
 from beeai_framework.memory import UnconstrainedMemory
 from beeai_framework.tools.search.wikipedia import WikipediaTool
 from beeai_framework.tools.weather.openmeteo import OpenMeteoTool
-from beeai_framework.workflows.agent import AgentFactoryInput, AgentWorkflow
+from beeai_framework.workflows.agent import AgentWorkflow
 
 
 async def main() -> None:
     llm = ChatModel.from_name("ollama:granite3.1-dense:8b")
 
     workflow = AgentWorkflow(name="Smart assistant")
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="Researcher",
-            instructions="You are a researcher assistant. Respond only if you can provide a useful answer.",
-            tools=[WikipediaTool()],
-            llm=llm,
-        )
+        name="Researcher",
+        instructions="You are a researcher assistant. Respond only if you can provide a useful answer.",
+        tools=[WikipediaTool()],
+        llm=llm,
     )
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="WeatherForecaster",
-            instructions="You are a weather assistant. Respond only if you can provide a useful answer.",
-            tools=[OpenMeteoTool()],
-            llm=llm,
-            execution=AgentExecutionConfig(max_iterations=3, total_max_retries=10, max_retries_per_step=3),
-        )
+        name="WeatherForecaster",
+        instructions="You are a weather assistant. Respond only if you can provide a useful answer.",
+        tools=[OpenMeteoTool()],
+        llm=llm,
+        execution=AgentExecutionConfig(max_iterations=3, total_max_retries=10, max_retries_per_step=3),
     )
     workflow.add_agent(
-        agent=AgentFactoryInput(
-            name="Solver",
-            instructions="""Your task is to provide the most useful final answer based on the assistants'
+        name="Solver",
+        instructions="""Your task is to provide the most useful final answer based on the assistants'
 responses which all are relevant. Ignore those where assistant do not know.""",
-            llm=llm,
-        )
+        llm=llm,
     )
 
     prompt = "What is the capital of France and what is the current weather there?"

diff --git a/python/tests/workflows/test_multi_agents.py b/python/tests/workflows/test_multi_agents.py
@@ -16,6 +16,7 @@
 
 from beeai_framework.adapters.ollama.backend.chat import OllamaChatModel
 from beeai_framework.agents.react import ReActAgent
+from beeai_framework.agents.types import AgentMeta
 from beeai_framework.backend.message import UserMessage
 from beeai_framework.memory import TokenMemory, UnconstrainedMemory
 from beeai_framework.workflows.agent import AgentFactoryInput, AgentWorkflow
@@ -31,7 +32,7 @@ async def test_multi_agents_workflow_basic() -> None:
     chat_model = OllamaChatModel()
 
     workflow: AgentWorkflow = AgentWorkflow()
-    workflow.add_agent(agent=AgentFactoryInput(name="Translator assistant", tools=[], llm=chat_model))
+    workflow.add_agent(AgentFactoryInput(name="Translator assistant", tools=[], llm=chat_model))
 
     memory = UnconstrainedMemory()
     await memory.add(UserMessage(content="Translate 'Hello' to German."))
@@ -47,7 +48,7 @@ async def test_multi_agents_workflow_creation() -> None:
 
     workflow: AgentWorkflow = AgentWorkflow()
     workflow.add_agent(ReActAgent(llm=chat_model, tools=[], memory=TokenMemory(chat_model)))
-    workflow.add_agent(agent=lambda mem: ReActAgent(llm=chat_model, tools=[], memory=mem))
+    workflow.add_agent(lambda mem: ReActAgent(llm=chat_model, tools=[], memory=mem))
 
     assert len(workflow.workflow.step_names) == 2
 
@@ -57,6 +58,40 @@ async def test_multi_agents_workflow_creation() -> None:
     assert "buongiorno" in response.state.final_answer.lower()
 
 
+@pytest.mark.e2e
+@pytest.mark.asyncio
+async def test_multi_agents_workflow_creation_variations() -> None:
+    chat_model = OllamaChatModel()
+
+    workflow: AgentWorkflow = AgentWorkflow()
+
+    # AgentFactoryInput
+    workflow.add_agent(name="AgentFactoryInput_1", tools=[], llm=chat_model)
+    workflow.add_agent(agent=AgentFactoryInput(name="AgentFactoryInput_2", tools=[], llm=chat_model))
+
+    # ReActAgent
+    workflow.add_agent(
+        agent=ReActAgent(
+            llm=chat_model,
+            tools=[],
+            memory=TokenMemory(chat_model),
+            meta=AgentMeta(name="ReActAgent_1", tools=[], description="ReActAgent defined using agent keyword"),
+        )
+    )
+
+    assert len(workflow.workflow.step_names) == 3
+    assert set(workflow.workflow.steps.keys()) == {
+        "AgentFactoryInput_1",
+        "AgentFactoryInput_2",
+        "ReActAgent_1",
+    }
+
+    memory = UnconstrainedMemory()
+    await memory.add(UserMessage(content="Translate 'Good morning' to Portuguese."))
+    response = await workflow.run(memory.messages)
+    assert "bom dia" in response.state.final_answer.lower()
+
+
 @pytest.mark.e2e
 @pytest.mark.asyncio
 async def test_multi_agents_workflow_agent_delete() -> None: