Use by code: Change to build and work with agents using code (#55)

* Removed passing config around for better code building apis

* Fixes after testing deletgator and reflection

* Setup notebook to demostrate composability by code

* Fix for bug in linear team team

* Renaming build team to flo

* Demostrating combining teams and agents in code

* RUnning delegator and reflection examples

Author: vizsatiz

examples/build_agents_by_code.ipynb

@@ -46,6 +46,6 @@
 )
 
 flo: Flo = Flo.build(session, yaml=yaml_data)
-flo.draw_to_file('delegate.png', xray=True)
-# data = flo.invoke(input_prompt)
-# print((data['messages'][-1]).content)
+Flo.set_log_level('INFO')
+data = flo.invoke(input_prompt)
+print((data['messages'][-1]).content)

examples/delegator_example.py

@@ -45,6 +45,5 @@
 )
 
 flo: Flo = Flo.build(session, yaml=yaml_data)
-flo.draw_to_file('event.png', xray=True)
 data = flo.invoke(input_prompt)
 print((data['messages'][-1]).content)

examples/email_reply_agent.ipynb

@@ -39,17 +39,17 @@ def parse_and_build_subteams(
     validate_team(name_set, team_config, session)
     if team_config.agents:
         members = [AgentFactory.create(session, agent) for agent in team_config.agents]
-        flo_team = FloTeam.Builder(team_config, members=members).build()
+        flo_team = FloTeam.Builder(session, team_config.name, members=members).build()
         router = FloRouterFactory.create(session, team_config, flo_team)
-        flo_routed_team = router.build_routed_team()
+        flo_routed_team = router.to_flo()
     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(team_config, members=flo_teams).build()
+        flo_team = FloTeam.Builder(session, team_config.name, members=flo_teams).build()
         router = FloRouterFactory.create(session, team_config, flo_team)
-        flo_routed_team = router.build_routed_team()
+        flo_routed_team = router.to_flo()
     return flo_routed_team
 
 

examples/reflection_example.py

@@ -7,6 +7,7 @@
 from flo_ai.models.flo_delegation_agent import FloDelegatorAgent
 from flo_ai.models.flo_tool_agent import FloToolAgent
 from flo_ai.error.flo_exception import FloException
+from flo_ai.models.delegate import Delegate
 from flo_ai.constants.common_constants import DOCUMENTATION_AGENT_ANCHOR
 from enum import Enum
 
@@ -61,8 +62,10 @@ def __create_agentic_agent(
         tools = [tool_map[tool.name] for tool in agent.tools]
         flo_agent: FloAgent = FloAgent.Builder(
             session,
-            agent,
-            tools,
+            name=agent.name,
+            job=agent.job,
+            tools=tools,
+            role=agent.role,
             llm=agent_model,
             on_error=session.on_agent_error,
             model_name=agent.model,
@@ -73,7 +76,12 @@ def __create_agentic_agent(
     def __create_llm_agent(session: FloSession, agent: AgentConfig) -> FloLLMAgent:
         agent_model = AgentFactory.__resolve_model(session, agent.model)
         builder = FloLLMAgent.Builder(
-            session, agent, llm=agent_model, model_name=agent.model
+            session,
+            name=agent.name,
+            job=agent.job,
+            role=agent.role,
+            llm=agent_model,
+            model_name=agent.model,
         )
         llm_agent: FloLLMAgent = builder.build()
         return llm_agent
@@ -82,19 +90,34 @@ def __create_llm_agent(session: FloSession, agent: AgentConfig) -> FloLLMAgent:
     def __create_runnable_agent(session: FloSession, agent: AgentConfig) -> FloLLMAgent:
         runnable = session.tools[agent.tools[0].name]
         return FloToolAgent.Builder(
-            session, agent, runnable, model_name=agent.model
+            session, agent.name, runnable, model_name=agent.model
         ).build()
 
     @staticmethod
     def __create_reflection_agent(
         session: FloSession, agent: AgentConfig
     ) -> FloReflectionAgent:
         agent_model = AgentFactory.__resolve_model(session, agent.model)
-        return FloReflectionAgent.Builder(session, agent, llm=agent_model).build()
+        return FloReflectionAgent.Builder(
+            session,
+            name=agent.name,
+            job=agent.job,
+            role=agent.role,
+            llm=agent_model,
+            to=Delegate([x.name for x in agent.to], agent.retry),
+            model_name=agent.model,
+        ).build()
 
     @staticmethod
     def __create_delegator_agent(
         session: FloSession, agent: AgentConfig
     ) -> FloReflectionAgent:
         agent_model = AgentFactory.__resolve_model(session, agent.model)
-        return FloDelegatorAgent.Builder(session, agent, llm=agent_model).build()
+        return FloDelegatorAgent.Builder(
+            session,
+            agent.name,
+            agent.job,
+            delegate=Delegate([x.name for x in agent.to], agent.retry),
+            llm=agent_model,
+            model_name=agent.model,
+        ).build()

flo_ai/builders/yaml_builder.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Delegate:
+    to: list[str]
+    retry: int = 1

flo_ai/factory/agent_factory.py

@@ -7,41 +7,39 @@
 from flo_ai.models.flo_executable import ExecutableFlo
 from flo_ai.state.flo_session import FloSession
 from typing import Union, Optional, Callable
-from flo_ai.yaml.config import AgentConfig
 from flo_ai.models.flo_executable import ExecutableType
 
 
 class FloAgent(ExecutableFlo):
     def __init__(
         self,
+        name: str,
         agent: Runnable,
         executor: AgentExecutor,
-        config: AgentConfig,
-        model_nick_name: str,
+        model_name: str,
     ) -> None:
-        super().__init__(config.name, executor, ExecutableType.agentic)
-        self.model_name = model_nick_name
+        super().__init__(name, executor, ExecutableType.agentic)
+        self.model_name = model_name
         self.agent: Runnable = (agent,)
         self.executor: AgentExecutor = executor
-        self.config: AgentConfig = config
 
     class Builder:
         def __init__(
             self,
             session: FloSession,
-            config: AgentConfig,
+            name: str,
+            job: str,
             tools: list[BaseTool],
-            verbose: bool = True,
             role: Optional[str] = None,
+            verbose: bool = True,
             llm: Union[BaseLanguageModel, None] = None,
             on_error: Union[str, Callable] = True,
             model_name: Union[str, None] = 'default',
         ) -> None:
-            prompt: Union[ChatPromptTemplate, str] = config.job
-            self.name: str = config.name
+            prompt: Union[ChatPromptTemplate, str] = job
+            self.name: str = name
             self.model_name = model_name
             self.llm = llm if llm is not None else session.llm
-            self.config = config
             system_prompts = (
                 [('system', 'You are a {}'.format(role)), ('system', prompt)]
                 if role is not None
@@ -67,6 +65,4 @@ def build(self) -> AgentExecutor:
                 return_intermediate_steps=True,
                 handle_parsing_errors=self.on_error,
             )
-            return FloAgent(
-                agent, executor, self.config, model_nick_name=self.model_name
-            )
+            return FloAgent(self.name, agent, executor, model_name=self.model_name)

flo_ai/models/delegate.py

@@ -1,40 +1,55 @@
 from typing import Optional
 from langchain_core.runnables import Runnable
-from flo_ai.yaml.config import AgentConfig
 from flo_ai.state.flo_session import FloSession
-from flo_ai.models.flo_executable import ExecutableFlo
+from flo_ai.models.flo_executable import ExecutableFlo, ExecutableType
+from flo_ai.models.delegate import Delegate
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from flo_ai.models.flo_executable import ExecutableType
-from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser
 from langchain_core.language_models import BaseLanguageModel
+from pydantic import BaseModel, Field
+from langchain_core.output_parsers import JsonOutputParser
+
+
+class NextAgent(BaseModel):
+    next: str = Field(description='Name of the next member to be called')
 
 
 class FloDelegatorAgent(ExecutableFlo):
     def __init__(
-        self, executor: Runnable, config: AgentConfig, model_name: str
+        self,
+        session: FloSession,
+        executor: Runnable,
+        delegate: Delegate,
+        name: str,
+        model_name: str,
     ) -> None:
-        super().__init__(config.name, executor, ExecutableType.delegator)
-        self.executor: Runnable = executor
-        self.config: AgentConfig = config
+        super().__init__(name, executor, ExecutableType.delegator)
+        self.session = session
+        self.delegate = delegate
+        self.executor = executor
         self.model_name = model_name
 
     class Builder:
         def __init__(
             self,
             session: FloSession,
-            agentConfig: AgentConfig,
+            name: str,
+            job: str,
+            delegate: Delegate,
             llm: Optional[BaseLanguageModel] = None,
             model_name: str = None,
         ) -> None:
-            self.config = agentConfig
+            self.session = session
+            self.name = name
+            self.to = delegate
             delegator_base_system_message = (
                 'You are a delegator tasked with routing a conversation between the'
                 ' following {member_type}: {members}. Given the following rules,'
                 ' respond with the worker to act next '
             )
             self.model_name = model_name
             self.llm = session.llm if llm is None else llm
-            self.options = [x.name for x in agentConfig.to]
+            self.options = delegate.to
+            self.parser = JsonOutputParser(pydantic_object=NextAgent)
             self.llm_router_prompt = ChatPromptTemplate.from_messages(
                 [
                     ('system', delegator_base_system_message),
@@ -43,43 +58,24 @@ def __init__(
                     (
                         'system',
                         'Given the conversation above, who should act next?'
-                        'Select one of: {options}',
+                        'Select one of: {options} \n {format_instructions}',
                     ),
                 ]
             ).partial(
                 options=str(self.options),
                 members=', '.join(self.options),
                 member_type='agents',
-                delegator_rules=agentConfig.job,
+                delegator_rules=job,
+                format_instructions=self.parser.get_format_instructions(),
             )
 
         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()
-            )
+            chain = self.llm_router_prompt | self.llm | self.parser
 
             return FloDelegatorAgent(
-                executor=chain, config=self.config, model_name=self.model_name
+                session=self.session,
+                name=self.name,
+                delegate=self.to,
+                executor=chain,
+                model_name=self.model_name,
             )

flo_ai/models/flo_agent.py

@@ -3,38 +3,36 @@
 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
+from typing import Union, Optional
 from langchain_core.output_parsers import StrOutputParser
-from flo_ai.yaml.config import AgentConfig
 from flo_ai.models.flo_executable import ExecutableType
 
 
 class FloLLMAgent(ExecutableFlo):
-    def __init__(
-        self, executor: Runnable, config: AgentConfig, model_name: str
-    ) -> None:
-        super().__init__(config.name, executor, ExecutableType.llm)
+    def __init__(self, name: str, executor: Runnable, model_name: str) -> None:
+        super().__init__(name, executor, ExecutableType.llm)
         self.executor: Runnable = executor
-        self.config: AgentConfig = config
         self.model_name: str = model_name
 
     class Builder:
         def __init__(
             self,
             session: FloSession,
-            config: AgentConfig,
+            name: str,
+            job: str,
+            role: Optional[str] = None,
             llm: Union[BaseLanguageModel, None] = None,
             model_name: str = None,
         ) -> None:
             self.model_name = model_name
-            prompt: Union[ChatPromptTemplate, str] = config.job
+            prompt: Union[ChatPromptTemplate, str] = job
 
-            self.name: str = config.name
+            self.name: str = name
             self.llm = llm if llm is not None else session.llm
             # TODO improve to add more context of what other agents are available
             system_prompts = (
-                [('system', 'You are a {}'.format(config.role)), ('system', prompt)]
-                if config.role is not None
+                [('system', 'You are a {}'.format(role)), ('system', prompt)]
+                if role is not None
                 else [('system', prompt)]
             )
             system_prompts.append(MessagesPlaceholder(variable_name='messages'))
@@ -43,8 +41,7 @@ def __init__(
                 if isinstance(prompt, str)
                 else prompt
             )
-            self.config = config
 
         def build(self) -> Runnable:
             executor = self.prompt | self.llm | StrOutputParser()
-            return FloLLMAgent(executor, self.config, self.model_name)
+            return FloLLMAgent(self.name, executor, self.model_name)