diff --git a/examples/linear_router_team.py b/examples/linear_router_team.py new file mode 100644 index 00000000..3ad5efa0 --- /dev/null +++ b/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) diff --git a/flo_ai/builders/yaml_builder.py b/flo_ai/builders/yaml_builder.py index 0dc49c50..14612b5b 100644 --- a/flo_ai/builders/yaml_builder.py +++ b/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, diff --git a/flo_ai/models/flo_team.py b/flo_ai/models/flo_team.py index 65053478..8c5fccce 100644 --- a/flo_ai/models/flo_team.py +++ b/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) diff --git a/flo_ai/router/flo_linear.py b/flo_ai/router/flo_linear.py index 2a2d0b23..93ae0839 100644 --- a/flo_ai/router/flo_linear.py +++ b/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" diff --git a/flo_ai/router/flo_router.py b/flo_ai/router/flo_router.py index 41d6bcf7..54a1453d 100644 --- a/flo_ai/router/flo_router.py +++ b/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