Fix for making outputs across agents and teams consistant

README.md

@@ -517,7 +517,7 @@ Built with ❤️ using:
         <p><b>Build an Agentic AI customer support bot using FloAI</b><br />We built an open-source agentic AI workflow builder named FloAI and used it to create an agentic customer support agent.</p>
     </a>
     <a href="https://medium.com/rootflo/build-an-agentic-rag-using-floai-in-minutes-0be260304c98" target="_blank" style="text-decoration: none;">
-        <img src="./examples/images/agentic-rag.png" width="150" style="border-radius: 10px;" />
+        <img src="./flo_ai/examples/images/agentic-rag.png" width="150" style="border-radius: 10px;" />
         <p><b>Build an Agentic RAG using FloAI in minutes</b><br />FloAI has just made implementing agentic RAG simple and easy to manage</p>
     </a>
 </div>

flo_ai/examples/email_reply_agent.ipynb

@@ -40,7 +40,7 @@
    "source": [
     "from typing import Optional, Type\n",
     "from pydantic import BaseModel, Field\n",
-    "from langchain.tools import BaseTool, StructuredTool, tool\n",
+    "from langchain.tools import BaseTool\n",
     "from langchain.callbacks.manager import (\n",
     "    AsyncCallbackManagerForToolRun,\n",
     "    CallbackManagerForToolRun,\n",

flo_ai/flo_ai/builders/yaml_builder.py

@@ -10,7 +10,6 @@
 from flo_ai.state.flo_session import FloSession
 from flo_ai.router.flo_router_factory import FloRouterFactory
 from flo_ai.factory.agent_factory import AgentFactory
-from flo_ai.error.flo_exception import FloException
 from flo_ai.yaml.validators import raise_for_name_error
 from flo_ai.common.flo_logger import get_logger
 
@@ -26,7 +25,7 @@ def build_supervised_team(
     elif isinstance(flo_config, FloAgentConfig):
         agent_config: AgentConfig = flo_config.agent
         validate_names(name_set, agent_config.name, session)
-        agent = AgentFactory.create(session, agent_config)
+        agent = parse_build_agent(session, agent_config, name_set)
         return agent
 
 
@@ -42,24 +41,34 @@ def parse_and_build_subteams(
     if team_config.agents:
         members = [AgentFactory.create(session, agent) for agent in team_config.agents]
         flo_team = FloTeam.Builder(session, team_config.name, members=members).build()
-        router = FloRouterFactory.create(session, team_config, flo_team)
+        router = FloRouterFactory.create(
+            session, team_config.router.kind, team_config, flo_team
+        )
         flo_routed_team = router.build_routed_team()
     else:
         flo_teams = []
         for subteam in team_config.subteams:
             flo_subteam = parse_and_build_subteams(session, subteam, name_set)
             flo_teams.append(flo_subteam)
         flo_team = FloTeam.Builder(session, team_config.name, members=flo_teams).build()
-        router = FloRouterFactory.create(session, team_config, flo_team)
+        router = FloRouterFactory.create(
+            session, team_config.router.kind, team_config, flo_team
+        )
         flo_routed_team = router.build_routed_team()
     return flo_routed_team
 
 
+def parse_build_agent(
+    session: FloSession, agent_config: AgentConfig, name_set=set()
+) -> ExecutableFlo:
+    validate_names(name_set, agent_config.name, session)
+    agent = AgentFactory.create(session, agent_config)
+    router = FloRouterFactory.create(session, 'agent', agent_config, agent)
+    return router.build_routed_team()
+
+
 def validate_names(name_set: set, name, session: FloSession):
     raise_for_name_error(name)
     if name in name_set:
-        get_logger().error(f"Duplicate name found: '{name}'", session)
-        raise FloException(
-            f"The name '{name}' is duplicate in the config. Make sure all teams and agents have unique names"
-        )
+        get_logger().warn(f"Duplicate name found: '{name}'", session)
     name_set.add(name)

flo_ai/flo_ai/core.py

@@ -8,7 +8,6 @@
 from typing import Any, Iterator, Union
 from flo_ai.router.flo_router import FloRouter
 from flo_ai.state.flo_session import FloSession
-from flo_ai.models.flo_llm_agent import FloLLMAgent
 from flo_ai.models.flo_executable import ExecutableFlo
 from flo_ai.error.flo_exception import FloException
 from flo_ai.constants.common_constants import DOCUMENTATION_WEBSITE
@@ -19,10 +18,10 @@
     set_logger_internal,
     FloLogConfig,
 )
-from langchain_core.messages import BaseMessage
-from flo_ai.models.flo_node import FloNode
 from flo_ai.models.flo_agent import FloAgent
+from flo_ai.models.flo_base_agent import FloBaseAgent
 from langchain.tools import StructuredTool
+from flo_ai.router.flo_agent_router import FloAgentRouter
 from flo_ai.callbacks.flo_execution_logger import ToolLogger
 
 
@@ -43,9 +42,7 @@ def async_stream(self, query, config=None) -> Iterator[Union[dict[str, Any], Any
         get_logger().info(f"Streaming async query requested: '{query}'", self.session)
         return self.runnable.astream(query, config)
 
-    def invoke(
-        self, query, config=None, chat_history: list[BaseMessage] = []
-    ) -> Iterator[Union[dict[str, Any], Any]]:
+    def invoke(self, query, config=None) -> Iterator[Union[dict[str, Any], Any]]:
         config = self.session.prepare_config(config)
 
         for callback in self.session.callbacks:
@@ -54,13 +51,11 @@ def invoke(
 
         self.validate_invoke(self.session)
         get_logger().info(f"Invoking query: '{query}'", self.session)
-        return self.runnable.invoke(query, config, chat_history)
+        return self.runnable.invoke(query, config)
 
-    def async_invoke(
-        self, query, config=None, chat_history: list[BaseMessage] = []
-    ) -> Iterator[Union[dict[str, Any], Any]]:
+    def async_invoke(self, query, config=None) -> Iterator[Union[dict[str, Any], Any]]:
         get_logger().info(f"Invoking async query: '{query}'", self.session)
-        return self.runnable.ainvoke(query, config, chat_history)
+        return self.runnable.ainvoke(query, config)
 
     @staticmethod
     def build(
@@ -96,9 +91,6 @@ def build(
             executable: ExecutableFlo = build_supervised_team(
                 session, to_supervised_team(yaml)
             )
-            # TODO fix this for all agents later
-            if isinstance(executable, FloAgent) or isinstance(executable, FloLLMAgent):
-                executable = FloNode.Builder(session).build_from_agent(executable)
             return Flo(session, executable)
         if routed_team is not None:
             return Flo(session, routed_team.build_routed_team())
@@ -108,8 +100,13 @@ def build(
     def create(session: FloSession, routed_team: Union[FloRouter, FloAgent]):
         if isinstance(routed_team, FloRouter):
             runnable = routed_team.build_routed_team()
-        else:
-            runnable = FloNode.Builder(session).build_from_agent(routed_team)
+        if isinstance(routed_team, FloBaseAgent):
+            agent_router = FloAgentRouter.Builder(
+                session,
+                f'router-{routed_team.name}',
+                flo_agent=routed_team,
+            ).build()
+            runnable = agent_router.build_routed_team()
         return Flo(session, runnable)
 
     @staticmethod

flo_ai/flo_ai/models/flo_agent.py

@@ -4,14 +4,15 @@
 from langchain_core.runnables import Runnable
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from flo_ai.models.flo_executable import ExecutableFlo, ExecutableType
+from flo_ai.models.flo_executable import ExecutableType
 from flo_ai.state.flo_session import FloSession
 from typing import Union, Optional, Callable
 from flo_ai.state.flo_output_collector import FloOutputCollector
 from flo_ai.parsers.flo_parser import FloParser
+from flo_ai.models.flo_base_agent import FloBaseAgent
 
 
-class FloAgent(ExecutableFlo):
+class FloAgent(FloBaseAgent):
     def __init__(
         self,
         name: str,
@@ -20,11 +21,15 @@ def __init__(
         model_name: str,
         data_collector: Optional[FloOutputCollector] = None,
     ) -> None:
-        super().__init__(name, executor, ExecutableType.agentic)
-        self.model_name = model_name
-        self.agent: Runnable = (agent,)
+        super().__init__(
+            name,
+            executor,
+            ExecutableType.agentic,
+            model_name,
+            data_collector=data_collector,
+        )
+        self.agent: Runnable = agent
         self.executor: AgentExecutor = executor
-        self.data_collector = data_collector
 
     @staticmethod
     def create(

flo_ai/flo_ai/models/flo_base_agent.py

@@ -0,0 +1,31 @@
+from typing import Optional
+from langchain_core.runnables import Runnable
+from flo_ai.models.flo_executable import ExecutableFlo, ExecutableType
+from flo_ai.state.flo_output_collector import FloOutputCollector
+
+
+class FloBaseAgent(ExecutableFlo):
+    """Base class for all Flo agents containing common properties and initialization."""
+
+    def __init__(
+        self,
+        name: str,
+        executor: Runnable,
+        executable_type: ExecutableType,
+        model_name: str,
+        data_collector: Optional[FloOutputCollector] = None,
+    ) -> None:
+        """Initialize the base agent with common properties.
+
+        Args:
+            name: Name of the agent
+            executor: The runnable executor for the agent
+            executable_type: Type of the executable
+            model_name: Name of the model being used
+            data_collector: Optional collector for output data
+        """
+        super().__init__(name, executor, executable_type)
+        self.executor: Runnable = executor
+        self.model_name: str = model_name
+        self.members = []
+        self.data_collector: Optional[FloOutputCollector] = data_collector

flo_ai/flo_ai/models/flo_delegation_agent.py

@@ -1,12 +1,13 @@
 from typing import Optional
 from langchain_core.runnables import Runnable
 from flo_ai.state.flo_session import FloSession
-from flo_ai.models.flo_executable import ExecutableFlo, ExecutableType
+from flo_ai.models.flo_executable import ExecutableType
 from flo_ai.models.delegate import Delegate
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.language_models import BaseLanguageModel
 from pydantic import BaseModel, Field
 from langchain_core.output_parsers import JsonOutputParser
+from flo_ai.models.flo_base_agent import FloBaseAgent
 
 
 # TODO probably use messages to relay information
@@ -15,7 +16,7 @@ class NextAgent(BaseModel):
     message: str = Field(description='Input to the next agent')
 
 
-class FloDelegatorAgent(ExecutableFlo):
+class FloDelegatorAgent(FloBaseAgent):
     def __init__(
         self,
         session: FloSession,
@@ -24,11 +25,9 @@ def __init__(
         name: str,
         model_name: str,
     ) -> None:
-        super().__init__(name, executor, ExecutableType.delegator)
+        super().__init__(name, executor, ExecutableType.delegator, model_name)
         self.session = session
         self.delegate = delegate
-        self.executor = executor
-        self.model_name = model_name
 
     @staticmethod
     def create(

flo_ai/flo_ai/models/flo_executable.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from flo_ai.models.flo_member import FloMember
 from langchain_core.runnables import Runnable
-from langchain_core.messages import HumanMessage, BaseMessage
+from langchain_core.messages import HumanMessage
 from flo_ai.state.flo_state import STATE_NAME_MESSAGES
 
 
@@ -36,15 +36,15 @@ def astream(self, work, config=None):
             {STATE_NAME_MESSAGES: [HumanMessage(content=work)]}, config
         )
 
-    def invoke(self, work, config=None, chat_history: list[BaseMessage] = []):
+    def invoke(self, work, config=None):
         return self.runnable.invoke(
             {
                 STATE_NAME_MESSAGES: ([HumanMessage(content=work)]),
             },
             config,
         )
 
-    def ainvoke(self, work, config=None, chat_history: list[BaseMessage] = []):
+    def ainvoke(self, work, config=None):
         return self.runnable.ainvoke(
             {
                 STATE_NAME_MESSAGES: [HumanMessage(content=work)],

flo_ai/flo_ai/models/flo_llm_agent.py

@@ -1,27 +1,30 @@
 from langchain_core.runnables import Runnable
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from flo_ai.models.flo_executable import ExecutableFlo
 from flo_ai.state.flo_session import FloSession
 from typing import Union, Optional
 from langchain_core.output_parsers import StrOutputParser
 from flo_ai.models.flo_executable import ExecutableType
 from flo_ai.parsers.flo_parser import FloParser
 from flo_ai.state.flo_output_collector import FloOutputCollector
+from flo_ai.models.flo_base_agent import FloBaseAgent
 
 
-class FloLLMAgent(ExecutableFlo):
+class FloLLMAgent(FloBaseAgent):
     def __init__(
         self,
         name: str,
         executor: Runnable,
         model_name: str,
         data_collector: Optional[FloOutputCollector] = None,
     ) -> None:
-        super().__init__(name, executor, ExecutableType.llm)
-        self.executor: Runnable = executor
-        self.model_name: str = model_name
-        self.data_collector = data_collector
+        super().__init__(
+            name,
+            executor,
+            ExecutableType.llm,
+            model_name,
+            data_collector=data_collector,
+        )
 
     @staticmethod
     def create(

flo_ai/flo_ai/models/flo_reflection_agent.py

@@ -1,20 +1,19 @@
 from typing import Union, Optional
 from langchain_core.runnables import Runnable
 from flo_ai.state.flo_session import FloSession
-from flo_ai.models.flo_executable import ExecutableFlo
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from flo_ai.models.flo_executable import ExecutableType
 from langchain_core.output_parsers import StrOutputParser
 from flo_ai.models.delegate import Delegate
+from flo_ai.models.flo_base_agent import FloBaseAgent
 
 
-class FloReflectionAgent(ExecutableFlo):
+class FloReflectionAgent(FloBaseAgent):
     def __init__(
         self, name: str, executor: Runnable, model_name: str, delegate: Delegate
     ) -> None:
-        super().__init__(name, executor, ExecutableType.reflection)
-        self.model_name = model_name
+        super().__init__(name, executor, ExecutableType.reflection, model_name)
         self.delegate = delegate
 
     @staticmethod

flo_ai/flo_ai/models/flo_tool_agent.py

@@ -1,23 +1,26 @@
 from typing import Optional
 from langchain_core.runnables import Runnable
-from flo_ai.models.flo_executable import ExecutableFlo
 from flo_ai.state.flo_session import FloSession
 from flo_ai.models.flo_executable import ExecutableType
 from flo_ai.state.flo_output_collector import FloOutputCollector
+from flo_ai.models.flo_base_agent import FloBaseAgent
 
 
-class FloToolAgent(ExecutableFlo):
+class FloToolAgent(FloBaseAgent):
     def __init__(
         self,
         name: str,
         executor: Runnable,
         model_name: str,
         data_collector: Optional[FloOutputCollector] = None,
     ) -> None:
-        super().__init__(name, executor, ExecutableType.tool)
-        self.executor: Runnable = executor
-        self.model_name: str = model_name
-        self.data_collector = data_collector
+        super().__init__(
+            name,
+            executor,
+            ExecutableType.tool,
+            model_name,
+            data_collector=data_collector,
+        )
 
     @staticmethod
     def create(