Router Pattern

from typing import Union,Literal,Optional,Dict,Any
from pydantic import BaseModel,Field

from stringdale import Define,Scope,V,E,Condition
from stringdale.chat import Chat

Sub Agents

rhyming_agent = Chat(model='gpt-4o-mini',
    messages=[{'role':'system','content':"""
    Answer the following questions using rhyming words.
    """},
    {'role':'user','content':'{{question}}'},
    ],
    )


await rhyming_agent(question='what is the capital of france?')

{'role': 'assistant',
 'content': "The capital of France, a city so bright,  \nIs Paris, where day turns to night.  \nWith its Eiffel Tower reaching high,  \nIt's a place where dreams seem to fly!",
 'meta': {'input_tokens': 28, 'output_tokens': 38}}

joke_agent = Chat(model='gpt-4o-mini',
    messages=[{'role':'system','content':"""
    Answer the following question with a joke.
    """},
    {'role':'user','content':'{{question}}'},
    ])

await joke_agent(question='what is the capital of france?')

{'role': 'assistant',
 'content': "Why did the Seine River break up with the Eiffel Tower? \n\nBecause it found out the tower was just a little too high-maintenance! \n\nAs for the capital of France, it's Paris!",
 'meta': {'input_tokens': 26, 'output_tokens': 39}}

yo_mama_chat = Chat(model='gpt-4o-mini',
    messages=[{'role':'system','content':"""
    Answer the following question with a joke about the person's mother.
    """},
    {'role':'user','content':'{{question}}'},
    ])

await yo_mama_chat(question='what is the capital of france?')

{'role': 'assistant',
 'content': 'Paris! But if your mom were answering, she\'d probably say, "The capital of France is \'ou la la!\' because she\'s always so dramatic!"',
 'meta': {'input_tokens': 30, 'output_tokens': 30}}

with Define('yo mama agent') as YoMamaAgent:
    V('yo mama',yo_mama_chat,inputs=['Start(**)'],outputs=['End'])
YoMamaAgent.draw()

yo_mama_agent = YoMamaAgent()
yo_mama_agent.run_all({
    'question':'what is the capital of france?'
    })

{'role': 'assistant',
 'content': 'Paris! But if your mom were answering, she\'d probably say, "The capital of France is \'ou la la!\' because she\'s always so dramatic!"',
 'meta': {'input_tokens': 30, 'output_tokens': 30}}

Router Factory

def RouterFactory(sub_agents:Dict[str,Any],default_choice:str):

    router = Chat(model='gpt-4o-mini',
        messages=[{'role':'system','content':"""
        Choose the best sub-agent to answer the following question from among the following options:
        {% for name,description in choice_descriptions.items() %}
        - {{name}}: {{description}}
        {% endfor %}
        """},
        {'role':'user','content':'{{question}}'},
        ],
        choice_descriptions = {name:agent['description'] for name,agent in sub_agents.items()},
        choices=list(sub_agents.keys())
    )

    with Define('Router',type='decision') as CompoundRouter:
        with Scope('flow'):
            V('choose_route',router,inputs=['Start(question=.)'])
            V('router',
                inputs=[
                    'choose_route(choice=content)',
                    'Start(question=.)'
                    ],)

        default = 'rhyme'
        for choice,agent in sub_agents.items():
            agent_func = agent['func']
            V(choice,agent_func,outputs=['End'])
            if choice == default_choice:
                E(f'router->{choice}(question=question)')
            else:
                E(f'router->{choice}(question=question)',cond=Condition(choice,'(0=choice)',name=f'choice=={choice}'))
    
    return CompoundRouter

sub_agents = {
    'rhyme':{
        'description':'this agent is good at rhyming',
        'func':rhyming_agent
    },
    'joke':{
        'description':'this agent is good at telling jokes',
        'func':joke_agent
    },
    'yo_mama':{
        'description':'this agent is specifically good at telling jokes about mothers',
        'func':yo_mama_chat
    }
}

Router = RouterFactory(sub_agents,'rhyme')
Router.draw()

d = Router()
for trace in d.run({'question':"what is the capital of france?, I like yo mama jokes"}):
    trace.pprint(skip_passthrough=True,drop_keys=['input'])

---
name: choose_route
output:
  role: assistant
  content: yo_mama
  meta:
    input_tokens: 208
    output_tokens: 11

---
name: yo_mama
output:
  role: assistant
  content: |-
    Why did your mother bring a ladder to the capital of France?   Because she heard
    the Eiffel Tower was a great place to elevate her sense of humor!
  meta:
    input_tokens: 40
    output_tokens: 31