Linear teams: Ability to execute teams using linear router

examples/linear_router_team.py

@@ -0,0 +1,53 @@
+from flo_ai import Flo
+from flo_ai.core import Flo
+from langchain_openai import ChatOpenAI
+from flo_ai import FloSession
+from langchain_community.tools.tavily_search.tool import TavilySearchResults
+from dotenv import load_dotenv
+load_dotenv()
+
+yaml_data = """
+apiVersion: flo/alpha-v1
+kind: FloRoutedTeam
+name: blogging-team
+team:
+    name: BloggingTeam
+    router:
+        name: linear-router
+        kind: linear
+    subteams:
+        - name: BloggingTeam
+          router:
+            name: supervisor
+            kind: supervisor
+          agents:
+            - name: Reasercher
+              job: Do a research on the internet and find articles of relevent to the topic asked by the user, always try to find the latest information on the same
+              tools:
+              - name: TavilySearchResults
+        - name: WritingTeam
+          router:
+            name: supervisor
+            kind: supervisor
+          agents: 
+            - name: Blogger
+              job: From the documents provider by the researcher write a blog of 300 words with can be readily published, make in engaging and add reference links to original blogs
+              tools:
+                - name: TavilySearchResults
+"""
+
+input_prompt = """
+Question: Write me an interesting blog about latest advancements in agentic AI
+"""
+
+
+llm = ChatOpenAI(temperature=0, model_name='gpt-4o')
+session = FloSession(llm).register_tool(
+    name="TavilySearchResults", tool=TavilySearchResults()
+)
+flo: Flo = Flo.build(session, yaml=yaml_data)
+
+for event in flo.stream(input_prompt):
+    for k, v in event.items():
+        if k != "__end__":
+            print(v)

flo_ai/builders/yaml_builder.py

@@ -43,7 +43,7 @@ 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 = FloRouterBuilder(session, team, agents).build()
+        router = FloRouterBuilder(session, team, flo_teams).build()
         flo_team = FloTeamBuilder(
             session=session,
             name=team.name,

flo_ai/models/flo_team.py

@@ -64,10 +64,13 @@ def build_node(self, flo_agent: FloAgent):
         agent_func = functools.partial(teamflo_agent_node, agent=flo_agent.executor, name=flo_agent.name)
         return FloAgentNode(agent_func, flo_agent.name)
 
-
+    # TODO re-structure to remove routing logic into corresponding routers and keep the team logic clean
     def build(self):
         if self.router_type == 'linear':
-            return self.build_linear_agents_graph()
+            if self.router.is_agent_supervisor():
+                return self.build_linear_agents_graph()
+            else:
+                return self.build_linear_teams_graph()
         elif self.router.is_agent_supervisor():
             return self.build_agent_supervisor_graph()
         else:
@@ -157,4 +160,31 @@ def build_linear_agents_graph(self) -> FloTeam:
         workflow_graph = workflow.compile()
 
         return FloTeam(self.name, workflow_graph)
+
+    def build_linear_teams_graph(self) -> FloTeam:
+        flo_team_entry_chains = [self.build_chain_for_teams(flo_agent) for flo_agent in self.flo_agents]
+        # Define the graph.
+        super_graph = StateGraph(TeamFloAgentState)
+        # First add the nodes, which will do the work
+        for flo_team_chain in flo_team_entry_chains:
+            agent_name = agent_name_from_randomized_name(flo_team_chain.name)
+            super_graph.add_node(agent_name, self.get_last_message | flo_team_chain.chain | self.join_graph)
 
+        if self.router.config.edges is None:
+            start_node_name = agent_name_from_randomized_name(flo_team_entry_chains[0].name)
+            end_node_name = agent_name_from_randomized_name(flo_team_entry_chains[-1].name)
+            super_graph.add_edge(START, start_node_name)
+            for i in range(len(flo_team_entry_chains) - 1):
+                agent1_name = agent_name_from_randomized_name(flo_team_entry_chains[i].name)
+                agent2_name = agent_name_from_randomized_name(flo_team_entry_chains[i+1].name)
+                super_graph.add_edge(agent1_name, agent2_name)
+            super_graph.add_edge(end_node_name, END)
+        else:
+            config = self.router.config
+            super_graph.add_edge(START, config.start_node)
+            for edge in config.edges:
+                super_graph.add_edge(edge[0], edge[1])
+            super_graph.add_edge(config.end_node, END)
+
+        super_graph = super_graph.compile()
+        return FloTeam(self.name, super_graph)

flo_ai/router/flo_linear.py

@@ -1,6 +1,12 @@
 from flo_ai.yaml.flo_team_builder import RouterConfig
+from flo_ai.models.flo_member import FloMember
 
 class FloLinear:
-    def __init__(self, agents, config):
+    def __init__(self, agents: list[FloMember], config: RouterConfig):
         self.agents = agents
+        self.name = config.name
         self.config: RouterConfig = config
+        self.type = agents[0].type
+
+    def is_agent_supervisor(self):
+        return self.type == "agent"

flo_ai/router/flo_router.py

@@ -4,7 +4,7 @@
 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
+from flo_ai.models.flo_member import FloMember
 
 class FloRouter:
     def __init__(self, kind, router):
@@ -16,7 +16,7 @@ class FloRouterBuilder:
     def __init__(self,
                 session: FloSession,
                 team_config: TeamConfig,
-                agents: list[FloAgent]
+                agents: list[FloMember]
             ) -> None:
         self.session = session
         self.team_config = team_config