LLM Router: A router that is controller by prompt

examples/linear_router_team.py

@@ -45,10 +45,10 @@ async def _arun(
           - 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: TavilySearchResults
           - 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:
+            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
 """
 

examples/llm_router_example.ipynb

@@ -0,0 +1,114 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from flo_ai import Flo\n",
+    "from flo_ai import FloSession\n",
+    "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n",
+    "\n",
+    "from dotenv import load_dotenv\n",
+    "load_dotenv()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain_community.tools.tavily_search.tool import TavilySearchResults\n",
+    "\n",
+    "llm = ChatOpenAI(temperature=0, model_name='gpt-4o-mini')\n",
+    "session = FloSession(llm)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "agent_yaml = \"\"\"\n",
+    "apiVersion: flo/alpha-v1\n",
+    "kind: FloRoutedTeam\n",
+    "name: adding-team\n",
+    "team:\n",
+    "    name: AddingTeam1\n",
+    "    router:\n",
+    "        name: router\n",
+    "        kind: llm\n",
+    "        job: Sent to the next member which was not called\n",
+    "    agents:\n",
+    "      - name: Agent1\n",
+    "        kind: llm\n",
+    "        role: Expert mathematician\n",
+    "        job: You are an expert in mathematics. Add one to the number given to you. And pass the result to the next agent\n",
+    "      - name: Agent2\n",
+    "        kind: llm\n",
+    "        role: Expert mathematician\n",
+    "        job: You are an expert in mathematics. Add one to the number given to you. And pass the result to the next agent\n",
+    "      - name: Agent3\n",
+    "        kind: llm\n",
+    "        role: Expert mathematician\n",
+    "        job: You are an expert in mathematics. Add one to the number given to you. And pass the result to the next agent\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'messages': [HumanMessage(content='Start with number 7'), HumanMessage(content='The result is 8.', name='Agent1-rzQ1B'), HumanMessage(content='The result is 9.', name='Agent2-izmIo'), HumanMessage(content='The result is 10.', name='Agent3-WgTbh')], 'next': 'FINISH'}\n"
+     ]
+    }
+   ],
+   "source": [
+    "flo = Flo.build(session, agent_yaml)\n",
+    "flo.draw_to_file(\"aasd.png\")\n",
+    "\n",
+    "print(flo.invoke(\"Start with number 7\"))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/planned_blogging_team.py

@@ -1,6 +1,5 @@
 from flo_ai.core import Flo
 from langchain_openai import ChatOpenAI
-from xamples.enviroment import load_env
 
 from flo_ai import FloSession
 from langchain_community.tools.tavily_search.tool import TavilySearchResults

flo_ai/__init__.py

@@ -1,6 +1,8 @@
 from flo_ai.core import Flo
 from flo_ai.models.flo_agent import FloAgent
-from flo_ai.router.flo_supervisor import FloSupervisorBuilder
-from flo_ai.models.flo_team import FloTeamBuilder
+from flo_ai.router.flo_supervisor import FloSupervisor
+from flo_ai.models.flo_team import FloTeam
+from flo_ai.router.flo_linear import FloLinear
 from flo_ai.state.flo_session import FloSession
 from flo_ai.retrievers.flo_retriever import FloRagBuilder
+

flo_ai/builders/yaml_builder.py

@@ -1,4 +1,4 @@
-from flo_ai.models.flo_team import FloTeamBuilder
+from flo_ai.models.flo_team import FloTeam
 from flo_ai.models.flo_agent import FloAgent
 from flo_ai.yaml.flo_team_builder import (FloRoutedTeamConfig, TeamConfig,
                                         AgentConfig, FloAgentConfig)
@@ -31,21 +31,19 @@ def parse_and_build_subteams(
         for agent in team_config.agents:
             flo_agent: FloAgent = AgentFactory.create(session, agent, tool_map)
             agents.append(flo_agent)
-        flo_team = FloTeamBuilder(
+        flo_team = FloTeam.Builder(
             session=session,
             name=team_config.name,
             members=agents
         ).build()
         router = FloRouterFactory.create(session, team_config, flo_team)
         flo_routed_team = router.build_routed_team()
-        if team_config.planner is not None:
-            return FloPlannerBuilder(session, team_config.planner.name, flo_routed_team).build()
     else:
         flo_teams = []
         for subteam in team_config.subteams:
             flo_subteam = parse_and_build_subteams(session, subteam, tool_map)
             flo_teams.append(flo_subteam)
-        flo_team = FloTeamBuilder(
+        flo_team = FloTeam.Builder(
             session=session,
             name=team_config.name,
             members=flo_teams

flo_ai/models/flo_node.py

@@ -18,7 +18,9 @@ class Builder():
         @staticmethod
         def teamflo_agent_node(state: TeamFloAgentState, agent: AgentExecutor, name: str):
             result = agent.invoke(state)
-            return { "messages": [HumanMessage(content=result["output"], name=name)] }
+            # TODO see how to fix this
+            output = result if isinstance(result, str) else result["output"]
+            return { "messages": [HumanMessage(content=output, name=name)] }
 
         @staticmethod
         def get_last_message(state: TeamFloAgentState) -> str:
@@ -29,7 +31,7 @@ def join_graph(response: dict):
             return { "messages": [ response["messages"][-1] ] }
 
         @staticmethod
-        def teamflo_team_node( message: str, members: list[str]):
+        def teamflo_team_node(message: str, members: list[str]):
             results = {
                 "messages": [HumanMessage(content=message)],
                 "team_members": ", ".join(members),
@@ -41,7 +43,7 @@ def build_from_agent(self, flo_agent: FloAgent):
             return FloNode(agent_func, flo_agent.name)
 
         def build_from_team(self, flo_team: FloRoutedTeam):
+            team_chain = (functools.partial(FloNode.Builder.teamflo_team_node, members=flo_team.runnable.nodes) | flo_team.runnable)
             return FloNode((
-                FloNode.Builder.get_last_message | functools.partial(FloNode.Builder.teamflo_team_node, members=flo_team.graph.nodes)
-                | flo_team.graph | FloNode.Builder.join_graph
+                FloNode.Builder.get_last_message | team_chain | FloNode.Builder.join_graph
             ), flo_team.name)

flo_ai/models/flo_team.py

@@ -12,19 +12,19 @@ def __init__(self,
         self.session = session
         self.members = members
 
-class FloTeamBuilder:
-    def __init__(self, 
-                 session: FloSession,
-                 name: str,
-                 members: list[FloMember]) -> None:
-        self.name = randomize_name(name)
-        self.session = session
-        self.members = members
-        self.member_names= list(map(lambda x: x.name, self.members))
-
-    def build(self):
-        return FloTeam(
-            name = self.name,
-            session=self.session, 
-            members=self.members
-        )
+    class Builder:
+        def __init__(self, 
+                    session: FloSession,
+                    name: str,
+                    members: list[FloMember]) -> None:
+            self.name = randomize_name(name)
+            self.session = session
+            self.members = members
+            self.member_names= list(map(lambda x: x.name, self.members))
+            
+        def build(self):
+            return FloTeam(
+                name = self.name,
+                session=self.session, 
+                members=self.members
+            )

flo_ai/router/flo_linear.py

@@ -67,4 +67,14 @@ def build_team_graph(self):
 
         super_graph = super_graph.compile()
         return FloRoutedTeam(self.flo_team.name, super_graph)
+
+    class Builder():
+
+        def __init__(self, session: FloSession, flo_team: FloTeam, config: RouterConfig) -> None:
+            self.config = config
+            self.session = session
+            self.team = flo_team
+
+        def build(self):
+            return FloLinear(self.session, self.team, self.config)
 

flo_ai/router/flo_llm_router.py

@@ -0,0 +1,134 @@
+from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from typing import Union
+from langchain_core.runnables import Runnable
+from flo_ai.state.flo_session import FloSession
+from flo_ai.constants.prompt_constants import FLO_FINISH
+from flo_ai.helpers.utils import randomize_name
+from flo_ai.router.flo_router import FloRouter
+from flo_ai.models.flo_team import FloTeam
+from flo_ai.models.flo_routed_team import FloRoutedTeam
+from langgraph.graph import StateGraph
+from flo_ai.state.flo_state import TeamFloAgentState
+
+class StateUpdateComponent:
+    def __init__(self, name: str, session: FloSession) -> None:
+        self.name = name
+        self.inner_session = session
+
+    def __call__(self, input):
+        self.inner_session.append(self.name)
+        return input
+
+class FloLLMRouter(FloRouter):
+
+    def __init__(self,
+                 session: FloSession,
+                 executor: Runnable, 
+                 flo_team: FloTeam,
+                 name: str) -> None:
+        super().__init__(
+            session = session, 
+            name = name, 
+            flo_team = flo_team,
+            executor = executor
+        )
+
+    def build_agent_graph(self):
+        flo_agent_nodes = [self.build_node(flo_agent) for flo_agent in self.members]
+        workflow = StateGraph(TeamFloAgentState)
+        for flo_agent_node in flo_agent_nodes:
+            workflow.add_node(flo_agent_node.name, flo_agent_node.func)
+        workflow.add_node(self.router_name, self.executor)
+        for member in self.member_names:
+            workflow.add_edge(member, self.router_name)
+        workflow.add_conditional_edges(self.router_name, self.router_fn)
+        workflow.set_entry_point(self.router_name)
+        workflow_graph = workflow.compile()
+        return FloRoutedTeam(self.flo_team.name, workflow_graph)
+
+    def build_team_graph(self):
+        flo_team_entry_chains = [self.build_node_for_teams(flo_agent) for flo_agent in self.members]
+        # 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:
+            super_graph.add_node(flo_team_chain.name, flo_team_chain.func)
+        super_graph.add_node(self.router_name, self.executor)
+
+        for member in self.member_names:
+            super_graph.add_edge(member, self.router_name)
+
+        super_graph.add_conditional_edges(self.router_name, self.router_fn)
+
+        super_graph.set_entry_point(self.router_name)
+        super_graph = super_graph.compile()
+        return FloRoutedTeam(self.flo_team.name, super_graph)
+
+    class Builder:
+        def __init__(self,
+                    session: FloSession,
+                    name: str,
+                    flo_team: FloTeam,
+                    router_prompt: ChatPromptTemplate = None,
+                    llm: Union[BaseLanguageModel, None] = None) -> None:
+
+            self.name = randomize_name(name)
+            self.session = session
+            self.llm = llm if llm is not None else session.llm
+            self.flo_team = flo_team
+            self.agents = flo_team.members
+            self.members = [agent.name for agent in flo_team.members]
+            self.options = self.members + [FLO_FINISH]
+            member_type = "workers" if flo_team.members[0].type == "agent" else "team members"
+
+            router_base_system_message = (
+                "You are a supervisor tasked with managing a conversation between the"
+                " following {member_type}: {members}. Given the following rules,"
+                " respond with the worker to act next "
+            )
+
+            self.llm_router_prompt = ChatPromptTemplate.from_messages(
+                [
+                    ("system", router_base_system_message),
+                    MessagesPlaceholder(variable_name="messages"),
+                    ("system", "Rules: {router_prompt}"),
+                    (
+                        "system",
+                        "Given the conversation above, who should act next?"
+                        " Or should we FINISH if the task is already answered ? Select one of: {options}",
+                    ),
+                ]
+            ).partial(options=str(self.options), members=", ".join(self.members), member_type=member_type, router_prompt=router_prompt)
+
+        def build(self):
+            function_def = {
+                "name": "route",
+                "description": "Select the next role.",
+                "parameters": {
+                    "title": "routeSchema",
+                    "type": "object",
+                    "properties": {
+                        "next": {
+                            "title": "Next",
+                            "anyOf": [
+                                {"enum": self.options},
+                            ],
+                        }
+                    },
+                    "required": ["next"],
+                }
+            }
+
+            chain = (
+                self.llm_router_prompt
+                | self.llm.bind_functions(functions=[function_def], function_call="route")
+                | JsonOutputFunctionsParser()
+                | StateUpdateComponent(self.name, self.session)
+            )
+
+            return FloLLMRouter(executor = chain, 
+                                flo_team=self.flo_team, 
+                                name=self.name, 
+                                session=self.session)