Introduce LLM routing & Linear Router Implementation

flo_ai/__init__.py

@@ -1,6 +1,7 @@
 from flo_ai.core import Flo
 from flo_ai.models.flo_agent import FloAgentBuilder
-from flo_ai.models.flo_supervisor import FloSupervisorBuilder
+from flo_ai.router.flo_supervisor import FloSupervisorBuilder
+from flo_ai.router.flo_router import FloRouterBuilder
 from flo_ai.models.flo_team import FloTeamBuilder
 from flo_ai.state.flo_session import FloSession
 from flo_ai.retrievers.flo_retriever import FloRagBuilder

flo_ai/builders/yaml_builder.py

@@ -1,12 +1,12 @@
 from flo_ai.models.flo_team import FloTeamBuilder
 from flo_ai.models.flo_agent import FloAgentBuilder, FloAgent
-from flo_ai.models.flo_supervisor import FloSupervisorBuilder
 from flo_ai.yaml.flo_team_builder import (FloRoutedTeamConfig, TeamConfig,
                                         AgentConfig, FloAgentConfig)
 from flo_ai.models.flo_executable import ExecutableFlo
 from flo_ai.models.flo_planner import FloPlannerBuilder
 from flo_ai.state.flo_session import FloSession
-from typing import Union, List
+from flo_ai.router.flo_router import FloRouterBuilder
+from typing import Union
 
 def build_supervised_team(
         session: FloSession,
@@ -30,11 +30,11 @@ def parse_and_build_subteams(
         for agent in team.agents:
             flo_agent: FloAgent = create_agent(session, agent, tool_map)
             agents.append(flo_agent)
-        router = create_router(session, team, agents)
+        router = FloRouterBuilder(session, team, agents).build()
         flo_team = FloTeamBuilder(
             session=session,
             name=team.name,
-            supervisor=router
+            router=router
         ).build()
         if team.planner is not None:
             return FloPlannerBuilder(session, team.planner.name, flo_team).build()
@@ -43,19 +43,13 @@ def parse_and_build_subteams(
         for subteam in team.subteams:
             flo_subteam = parse_and_build_subteams(session, subteam, tool_map)
             flo_teams.append(flo_subteam)
-        router = create_router(session, team, agents)
+        router = FloRouterBuilder(session, team, agents).build()
         flo_team = FloTeamBuilder(
             session=session,
             name=team.name,
-            supervisor=router
+            router=router
         ).build()
     return flo_team
-
-def create_router(session: FloSession, team_config: TeamConfig, agents: List[FloAgent]):
-    if (team_config.router.kind == "supervisor"):
-            return FloSupervisorBuilder(session, team_config.router.name, agents).build()
-    else:
-        raise Exception("Unknown router type")
 
 def create_agent(session: FloSession, agent: AgentConfig, tool_map) -> FloAgent:
     tools = [tool_map[tool.name] for tool in agent.tools]

flo_ai/helpers/utils.py

@@ -7,4 +7,7 @@ def random_str(length: int = 5):
     return result_str
 
 def randomize_name(name: str):
-    return "{}-{}".format(name, random_str(5))
+    return "{}-{}".format(name, random_str(5))
+
+def agent_name_from_randomized_name(name: str):
+    return name.split("-")[0]

flo_ai/models/flo_team.py

@@ -1,16 +1,19 @@
 import functools
 from langchain_core.messages import HumanMessage
 from langgraph.graph.graph import CompiledGraph
-from langgraph.graph import StateGraph, END
+from langgraph.graph import StateGraph, END, START
 from langchain.agents import AgentExecutor
 
 from flo_ai.models.flo_executable import ExecutableFlo
 from flo_ai.models.flo_agent import FloAgent
-from flo_ai.models.flo_supervisor import FloSupervisor
+from flo_ai.router.flo_supervisor import FloSupervisor
+from flo_ai.router.flo_linear import FloLinear
 from flo_ai.constants.prompt_constants import FLO_FINISH
 from flo_ai.state.flo_state import TeamFloAgentState
 from flo_ai.state.flo_session import FloSession
-from flo_ai.helpers.utils import randomize_name
+from flo_ai.helpers.utils import randomize_name, agent_name_from_randomized_name
+from flo_ai.router.flo_router import FloRouter
+from typing import Union
 
 class FloAgentNode:
     def __init__(self, 
@@ -49,11 +52,12 @@ class FloTeamBuilder:
     def __init__(self, 
                  session: FloSession,
                  name: str,
-                 supervisor: FloSupervisor) -> None:
+                 router: FloRouter) -> None:
         self.name = randomize_name(name)
         self.session = session
-        self.flo_agents = supervisor.agents
-        self.flo_supervisor: FloSupervisor = supervisor
+        self.flo_agents = router.output.agents
+        self.router_type = router.kind
+        self.router: Union[FloSupervisor, FloLinear] = router.output
         self.members = list(map(lambda x: x.name, self.flo_agents))
 
     def build_node(self, flo_agent: FloAgent):
@@ -62,12 +66,14 @@ def build_node(self, flo_agent: FloAgent):
 
 
     def build(self):
-        if self.flo_supervisor.is_agent_supervisor():
+        if self.router_type == 'linear':
+            return self.build_linear_agents_graph()
+        elif self.router.is_agent_supervisor():
             return self.build_agent_supervisor_graph()
         else:
             return self.build_team_supervisor_graph()
 
-    def router(self, state: TeamFloAgentState):
+    def router_fn(self, state: TeamFloAgentState):
         next = state["next"]
         conditional_map = {k: k for k in self.members}
         conditional_map[FLO_FINISH] = END
@@ -82,13 +88,13 @@ def build_agent_supervisor_graph(self) -> FloTeam:
         workflow = StateGraph(TeamFloAgentState)
         for flo_agent_node in flo_agent_nodes:
             workflow.add_node(flo_agent_node.name, flo_agent_node.agent_node)
-        workflow.add_node(self.flo_supervisor.name, self.flo_supervisor.executor)
+        workflow.add_node(self.router.name, self.router.executor)
         for member in self.members:
-            workflow.add_edge(member, self.flo_supervisor.name)
+            workflow.add_edge(member, self.router.name)
 
-        workflow.add_conditional_edges(self.flo_supervisor.name, self.router)
+        workflow.add_conditional_edges(self.router.name, self.router_fn)
 
-        workflow.set_entry_point(self.flo_supervisor.name)
+        workflow.set_entry_point(self.router.name)
         workflow_graph = workflow.compile()
         return FloTeam(self.name, workflow_graph)
 
@@ -113,14 +119,42 @@ def build_team_supervisor_graph(self) -> FloTeam:
         # First add the nodes, which will do the work
         for flo_team_chain in flo_team_entry_chains:
             super_graph.add_node(flo_team_chain.name, self.get_last_message | flo_team_chain.chain | self.join_graph)
-        super_graph.add_node(self.flo_supervisor.name, self.flo_supervisor.executor)
+        super_graph.add_node(self.router.name, self.router.executor)
 
         for member in self.members:
-            super_graph.add_edge(member, self.flo_supervisor.name)
+            super_graph.add_edge(member, self.router.name)
 
-        super_graph.add_conditional_edges(self.flo_supervisor.name, self.router)
+        super_graph.add_conditional_edges(self.router.name, self.router_fn)
 
-        super_graph.set_entry_point(self.flo_supervisor.name)
+        super_graph.set_entry_point(self.router.name)
         super_graph = super_graph.compile()
         return FloTeam(self.name, super_graph)
+
+    def build_linear_agents_graph(self) -> FloTeam:
+        flo_agent_nodes = [self.build_node(flo_agent) for flo_agent in self.flo_agents]
+        workflow = StateGraph(TeamFloAgentState)
+
+        for flo_agent_node in flo_agent_nodes:
+            agent_name = agent_name_from_randomized_name(flo_agent_node.name)
+            workflow.add_node(agent_name, flo_agent_node.agent_node)
+
+        if self.router.config.edges is None:
+            start_node_name = agent_name_from_randomized_name(flo_agent_nodes[0].name)
+            end_node_name = agent_name_from_randomized_name(flo_agent_nodes[-1].name)
+            workflow.add_edge(START, start_node_name)
+            for i in range(len(flo_agent_nodes) - 1):
+                agent1_name = agent_name_from_randomized_name(flo_agent_nodes[i].name)
+                agent2_name = agent_name_from_randomized_name(flo_agent_nodes[i+1].name)
+                workflow.add_edge(agent1_name, agent2_name)
+            workflow.add_edge(end_node_name, END)
+        else:
+            config = self.router.config
+            workflow.add_edge(START, config.start_node)
+            for edge in config.edges:
+                workflow.add_edge(edge[0], edge[1])
+            workflow.add_edge(config.end_node, END)
+
+        workflow_graph = workflow.compile()
+
+        return FloTeam(self.name, workflow_graph)
 

flo_ai/router/flo_linear.py

@@ -0,0 +1,6 @@
+from flo_ai.yaml.flo_team_builder import RouterConfig
+
+class FloLinear:
+    def __init__(self, agents, config):
+        self.agents = agents
+        self.config: RouterConfig = config

flo_ai/router/flo_router.py

@@ -0,0 +1,34 @@
+
+from typing import Union, List
+from flo_ai.state.flo_session import FloSession
+from flo_ai.router.flo_supervisor import FloSupervisorBuilder, FloSupervisor
+from flo_ai.router.flo_linear import FloLinear
+from flo_ai.yaml.flo_team_builder import TeamConfig
+from flo_ai.models.flo_agent import  FloAgent
+
+class FloRouter:
+    def __init__(self, kind, router):
+        self.kind: str = kind
+        self.output: Union[FloSupervisor, FloLinear] = router
+
+
+class FloRouterBuilder:
+    def __init__(self,
+                session: FloSession,
+                team_config: TeamConfig,
+                agents: list[FloAgent]
+            ) -> None:
+        self.session = session
+        self.team_config = team_config
+        self.agents = agents
+
+
+    def build(self) -> FloRouter:
+        if self.team_config.router.kind == 'supervisor':
+            router = FloSupervisorBuilder(self.session, self.team_config.router, self.agents).build()
+        elif self.team_config.router.kind == 'linear':
+            router = FloLinear(self.agents, self.team_config.router)
+        else:
+            raise Exception("Unknown router type")
+
+        return FloRouter(self.team_config.router.kind, router)

flo_ai/yaml/flo_team_builder.py

@@ -39,6 +39,9 @@ class AgentConfig(BaseModel):
 class RouterConfig(BaseModel):
     name: str
     kind: str
+    start_node: Optional[str] = None
+    end_node: Optional[str] = None
+    edges: Optional[List[List[str]]] = None
 
 class PlannerConfig(BaseModel):
     name: str