Model Switching

Use multiple models within the same workflow

The latest version of flo-ai (>= 4.0.0) lets you use multiple models within the same workflow. Now you can specify the model to be used by each agent.

This involves following steps:

1. Registering the models with the FloSession

2. Using the model name in the yaml

When no exclusive model is specified, the system will fall back to model defined within the session, and call it the default model

Using 3 models with same workflow

Create your model objects.

from flo_ai import Flo
from flo_ai import FloSession
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI

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)

Register all the models in the session

session.register_model('bronze', gpt35)
session.register_model('silver', gpt_4o_mini)
session.register_model('gold', gpt_4o)

Define the yaml schema as use the model name registered

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
"""

This lets you use different models in different agents, here is how this works:

1. EmailFriends: This agent will use the bronze model, which is in this case gpt35

2. EmailColleagues: This agent will use the gold model, which in this case is gpt4o

3. The router Personal-Assistant will use silver model to work with these agents

If you are using code, you can directly pass the model object like below:

blogger = FloAgent.create(
    session, 
    llm=<llm object>
    name="BlogWriter", 
    role="Thought Leader", # optional
    job="Able to write a blog using information provided", 
    tools=[TavilySearchResults()]
)