Adding the ability to use different models for different components

.github/workflows/build-project.yml

@@ -31,3 +31,6 @@ jobs:
 
     - name: Run build
       run: poetry build
+
+    - name: Run tests
+      run: poetry run pytest

ROADMAP.md

@@ -10,8 +10,9 @@ Core features improve the library itself to cater wider range of functionalities
 |------|-------------|--------|-----------------|
 | Full composability | Right now teams can only be combined with teams and agents with agents. We want to extend this to team + agent composibility | In progress | 0.0.4 | 
 | Error handling | Ability to handle errors autonomously | In Progress | 0.0.4|
+|LLM Extensibilty| Ability to different LLMs across different agents and teams| In Progress | 0.0.4|
+|Async Tools| Ability create tools easily within asyncio | In Progress | 0.0.4|
 |Output formatter| Ability to templatize output format using pydantic| Yet to start| 0.0.4|
-|LLM Extensibilty| Ability to different LLMs across different agents and teams| Yet to start| 0.0.4|
 |Auto-Query RAG| Ability to make metadata query within the agentic, which can automatically add metadata while rag query runs, like timestamp or other data|Yet to start|TBD|
 |Parellel Router| A router to execute tasks or agents in parallel if the tasks are independent | Yet to start | TBD
 

examples/hierarchical_blogging_team.py

@@ -49,4 +49,4 @@
 session = FloSession(llm).register_tool(
     name="TavilySearchResults", tool=TavilySearchResults()
 )
-flo: Flo = Flo.build(session, yaml=yaml_data)
+flo: Flo = Flo.build(session, yaml=yaml_data)

examples/linear_router_example.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -19,7 +19,7 @@
    "source": [
     "from flo_ai import Flo\n",
     "from flo_ai import FloSession\n",
-    "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n",
+    "from langchain_openai import ChatOpenAI\n",
     "\n",
     "from dotenv import load_dotenv\n",
     "load_dotenv()"

examples/llm_extensibility.py

@@ -0,0 +1,79 @@
+import os
+from flo_ai import Flo
+from flo_ai import FloSession
+from pydantic import BaseModel, Field
+from langchain_openai import ChatOpenAI, AzureChatOpenAI
+from flo_ai.tools.flo_tool import flotool
+
+from dotenv import load_dotenv
+load_dotenv()
+
+import warnings
+warnings.simplefilter('default', DeprecationWarning)
+
+gpt35 = ChatOpenAI(temperature=0, model_name='gpt-3.5-turbo')
+gpt_4o_mini = ChatOpenAI(temperature=0, model_name='gpt-4o-mini')
+gpt_4o = ChatOpenAI(temperature=0, model_name='gpt-4o')
+session = FloSession(gpt35)
+
+session.register_model("bronze", gpt35)
+session.register_model("silver", gpt_4o_mini)
+session.register_model("gold", gpt_4o)
+
+class SendEmailInput(BaseModel):
+    to: str = Field(description="Comma seperared list of users emails to which email needs to be sent")
+    message: str = Field(description="The email text to be sent")
+
+@flotool("email_triage", "useful for when you need to send an email to someone", argument_contract=SendEmailInput)
+def email_tool(to: str, message: str):
+    return f"Email sent successfully to: {to}"
+
+session.register_tool("SendEmailTool", email_tool)
+
+agent_yaml = """
+apiVersion: flo/alpha-v1
+kind: FloRoutedTeam
+name: invite-handler
+team:
+    name: Personal-Assistant-Bot
+    router:
+        name: Personal-Assistant
+        kind: supervisor
+        model: silver
+    agents:
+      - name: EmailFriends
+        job: You job is to send an invite to the christmas party at my house to my friends and friends only, not collegues, invite their spouses too. Keep the email warm and friendly.
+        role: personal ai assistant
+        model: bronze
+        tools:
+          - name: SendEmailTool
+      - name: EmailColleagues
+        job: You job is to send an invite to the christmas party at my house to my colleagues and not friends. Keep the email formal, and DO NOT invite the spouses.
+        role: office ai assistant
+        model: gold
+        tools:
+          - name: SendEmailTool
+"""
+input_prompt = """
+Here is the list of user emails and there relations to me
+
+vishnu@gmail.com / friend
+nk@gmail.com / friend
+jk@gmail.com / colleague
+ck@hotmail.com / friend
+hk@gmail.com / colleague
+jak@gmail.com / colleague
+ck@gmail.com / friend.
+
+Please invite these nice folks to my christmas party
+"""
+
+flo: Flo = Flo.build(session, yaml=agent_yaml)
+for s in flo.stream(input_prompt):
+     if "__end__" not in s:
+        print(s)
+        print("----")
+
+
+
+

examples/tool_error_handling.py

@@ -1,7 +1,6 @@
 from flo_ai import Flo
 from flo_ai import FloSession
 from langchain_openai import ChatOpenAI
-from langchain.tools import BaseTool
 from pydantic import BaseModel, Field
 from typing import List
 from dotenv import load_dotenv
@@ -67,6 +66,7 @@ async def invoke_main():
 
 asyncio.run(invoke_main())
 
+
 # import asyncio
 
 # async def stream_main():

flo_ai/builders/yaml_builder.py

@@ -46,5 +46,5 @@ def validate_names(name_set: set, name):
     raise_for_name_error(name)
     if name in name_set:
         builder_logger.error(f"Duplicate name found: '{name}'")
-        raise FloException(f"The name '{name}' is already in the set.")
+        raise FloException(f"The name '{name}' is duplicate in the config. Make sure all teams and agents have unique names")
     name_set.add(name)

flo_ai/constants/common_constants.py

@@ -0,0 +1,4 @@
+DOCUMENTATION_WEBSITE = "https://flo-ai.rootflo.ai"
+DOCUMENTATION_ROUTER_ANCHOR = f"{DOCUMENTATION_WEBSITE}/basics/routers"
+DOCUMENTATION_AGENT_ANCHOR = f"{DOCUMENTATION_WEBSITE}/basics/agents"
+DOCUMENTATION_AGENT_TOOLS = f"{DOCUMENTATION_WEBSITE}/basics/tools"

flo_ai/core.py

@@ -5,6 +5,7 @@
 from flo_ai.state.flo_session import FloSession
 from flo_ai.models.flo_executable import ExecutableFlo
 from flo_ai.error.flo_exception import FloException
+from flo_ai.constants.common_constants import DOCUMENTATION_WEBSITE
 from flo_ai.common.flo_logger import common_logger, builder_logger, set_global_log_level
 
 class Flo:
@@ -24,6 +25,7 @@ def __init__(self,
         self.logger.info(f"Flo instance created for session {session.session_id}")
 
     def stream(self, query, config = None) -> Iterator[Union[dict[str, Any], Any]]:
+        self.validate_invoke(self.session)
         self.logger.info(f"Streaming query for session {self.session.session_id}: {query}")
         return self.runnable.stream(query, config)
 
@@ -35,6 +37,7 @@ def invoke(self, query, config = None) -> Iterator[Union[dict[str, Any], Any]]:
         config = {
          'callbacks' : [self.session.langchain_logger]
         }
+        self.validate_invoke(self.session)
         self.logger.info(f"Invoking query for session {self.session.session_id}: {query}")
         return self.runnable.invoke(query, config)
 
@@ -63,4 +66,8 @@ def draw_to_file(self, filename: str, xray=True):
             image = PILImage.open(image_io)
             image.save(filename)
 
+    def validate_invoke(self, session: FloSession):
+        async_coroutines = filter(lambda x: (hasattr(x, "coroutine") and asyncio.iscoroutinefunction(x.coroutine)) ,session.tools.values())
+        if len(list(async_coroutines)) > 0:
+            raise FloException(f"""You seem to have atleast one async tool registered in this session. Please use flo.async_invoke or flo.async_stream. Checkout {DOCUMENTATION_WEBSITE}""")
 

flo_ai/factory/agent_factory.py

@@ -1,11 +1,13 @@
+from typing import Optional
 from flo_ai.state.flo_session import FloSession
 from flo_ai.yaml.config import (AgentConfig)
 from flo_ai.models.flo_agent import FloAgent
 from flo_ai.models.flo_llm_agent import FloLLMAgent
 from flo_ai.models.flo_reflection_agent import FloReflectionAgent
 from flo_ai.models.flo_delegation_agent import FloDelegatorAgent
 from flo_ai.models.flo_tool_agent import FloToolAgent
-from flo_ai.models.flo_executable import ExecutableFlo, ExecutableType
+from flo_ai.error.flo_exception import FloException
+from flo_ai.constants.common_constants import DOCUMENTATION_AGENT_ANCHOR
 from enum import Enum
 
 class AgentKinds(Enum):
@@ -25,7 +27,8 @@ def create(session: FloSession, agent: AgentConfig):
         if kind is not None:
             agent_kind = getattr(AgentKinds, kind, None)
             if agent_kind is None:
-                raise ValueError(f"Agent kind cannot be {kind} !")
+                raise ValueError(f"""Unknown agent kind: `{kind}`. The supported types are llm, tool, reflection, delegator or agentic. 
+                            Check the documentation @ {DOCUMENTATION_AGENT_ANCHOR}""")
             match(agent_kind):
                 case AgentKinds.llm:
                     return AgentFactory.__create_llm_agent(session, agent)
@@ -36,21 +39,35 @@ def create(session: FloSession, agent: AgentConfig):
                 case AgentKinds.delegator:
                     return AgentFactory.__create_delegator_agent(session, agent)
         return AgentFactory.__create_agentic_agent(session, agent, tool_map)
+
+    @staticmethod
+    def __resolve_model(session: FloSession, model_name: Optional[str] = None):
+        if model_name is None:
+            return session.llm
+        if model_name not in session.models:
+            raise FloException(
+                f"""Model not found: {model_name}.
+                The model you would like to use should be registered to the session using session.register_model api, 
+                and the same model name should be used here instead of `{model_name}`""")
+        return session.models[model_name]
 
     @staticmethod
     def __create_agentic_agent(session: FloSession, agent: AgentConfig, tool_map) -> FloAgent:
+        agent_model = AgentFactory.__resolve_model(session, agent.model)
         tools = [tool_map[tool.name] for tool in agent.tools]
         flo_agent: FloAgent = FloAgent.Builder(
             session,
             agent,
             tools,
+            llm=agent_model,
             on_error=session.on_agent_error
         ).build()
         return flo_agent
 
     @staticmethod
     def __create_llm_agent(session: FloSession, agent: AgentConfig) -> FloLLMAgent:
-        builder = FloLLMAgent.Builder(session, agent)
+        agent_model = AgentFactory.__resolve_model(session, agent.model)
+        builder = FloLLMAgent.Builder(session, agent, llm=agent_model)
         llm_agent: FloLLMAgent = builder.build()
         return llm_agent
 
@@ -61,8 +78,10 @@ def __create_runnable_agent(session: FloSession, agent: AgentConfig) -> FloLLMAg
 
     @staticmethod
     def __create_reflection_agent(session: FloSession, agent: AgentConfig) -> FloReflectionAgent:
-        return FloReflectionAgent.Builder(session, agent).build()
+        agent_model = AgentFactory.__resolve_model(session, agent.model)
+        return FloReflectionAgent.Builder(session, agent, llm=agent_model).build()
 
     @staticmethod
     def __create_delegator_agent(session: FloSession, agent: AgentConfig) -> FloReflectionAgent:
-        return FloDelegatorAgent.Builder(session, agent).build()
+        agent_model = AgentFactory.__resolve_model(session, agent.model)
+        return FloDelegatorAgent.Builder(session, agent, llm=agent_model).build()

flo_ai/models/flo_delegation_agent.py

@@ -1,10 +1,12 @@
+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 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
 
 
 class FloDelegatorAgent(ExecutableFlo):
@@ -20,14 +22,15 @@ def __init__(self,
     class Builder():
         def __init__(self, 
                      session: FloSession,
-                     agentConfig: AgentConfig) -> None:
+                     agentConfig: AgentConfig,
+                     llm: Optional[BaseLanguageModel] = None) -> None:
             self.config = agentConfig
             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.llm = session.llm
+            self.llm = session.llm if llm is None else llm
             self.options = [x.name for x in agentConfig.to]
             self.llm_router_prompt = ChatPromptTemplate.from_messages(
                 [

flo_ai/models/flo_tool_agent.py

@@ -22,6 +22,5 @@ def __init__(self,
             self.runnable = tool_runnable
             self.config = config
 
-
         def build(self) -> Runnable:
             return FloToolAgent(self.runnable, self.config)