BrainBlend-AI
Orchestrator example calling agents and tools
Hi and thanks for this library! I am trying to understand how it works, what should I call etc..
I will be very grateful if someone can guide me through the following example.
What I want to do is to have an orchestrator agent which will decide which agent to call. Each agent should have access to specific tools (functions).
-
modification agent: modify the csv data
- tools : modify_weather_data
-
forecast agent: makes the forecast
- tools: run_forecast_model
-
analysis agent: runs the analysis and provides results (plots, metrics)
- tools: create_comparison_plot , calculate_impact_metrics, format_impact_analysis
(the tools are just python functions. I am not sure if we have to create custom Tools??)
So, the setup I use:
# ─── 1) Modification Data Agent
class ModificationInputSchema(BaseIOSchema):
'''Modify data'''
data_path: str = Field(..., description="Path to original weather data CSV")
variable: str = Field(..., description="Weather variable to perturb")
scenario_type: str = Field(..., description="Scenario template text")
value: float = Field(..., description="Perturbation magnitude")
days: int | None = Field(None, description="Duration in days, if applicable")
class ModificationOutputSchema(BaseIOSchema):
'''Modified data'''
modified_data_path: str = Field(..., description="Path to the modified data CSV")
scenario_desc: str = Field(..., description="Description of modification applied")
mod_system_prompt = SystemPromptGenerator(
background=["You are a data analysis agent."],
steps=[
"Validate parameters",
"Decide call to modify_weather_data tool",
"Emit JSON matching ModificationOutputSchema"
],
output_instructions=["Return JSON exactly matching ModificationOutputSchema"]
)
modification_agent = BaseAgent(
config=BaseAgentConfig(
client=client,
model=model,
system_prompt_generator=mod_system_prompt,
memory=AgentMemory(),
input_schema=ModificationInputSchema,
output_schema=ModificationOutputSchema
)
)
# ─── 2) Forecast Agent
class ForecastInputSchema(BaseIOSchema):
'''Forecast'''
data_path: str = Field(..., description="CSV for forecasting")
variable: str = Field(..., description="Target variable to forecast")
periods: int = Field(..., description="Days ahead to forecast")
is_original: bool = Field(..., description="True for original data, False for modified")
class ForecastOutputSchema(BaseIOSchema):
'''Forecast'''
forecast_path: str = Field(..., description="Path to the forecast CSV")
forecast_system_prompt = SystemPromptGenerator(
background=["You are a forecasting agent."],
steps=[
"Validate inputs",
"Decide call to run_forecast_model tool",
"Emit JSON matching ForecastOutputSchema"
],
output_instructions=["Return JSON exactly matching ForecastOutputSchema"]
)
forecast_agent = BaseAgent(
config=BaseAgentConfig(
client=client,
model=model,
system_prompt_generator=forecast_system_prompt,
memory=AgentMemory(),
input_schema=ForecastInputSchema,
output_schema=ForecastOutputSchema
)
)
# ─── 3) Analysis Agent
class AnalysisInputSchema(BaseIOSchema):
'''Analysis'''
original_forecast_path: str = Field(..., description="Original forecast CSV path")
modified_forecast_path: str = Field(..., description="Modified forecast CSV path")
variable: str = Field(..., description="Variable being compared")
class AnalysisOutputSchema(BaseIOSchema):
'''Analysis'''
comparison_plot_path: str = Field(..., description="Path to saved comparison plot")
impact_metrics: dict = Field(..., description="Raw metrics dictionary")
impact_analysis: str = Field(..., description="Formatted markdown analysis")
analysis_system_prompt = SystemPromptGenerator(
background=["You are an analysis agent."],
steps=[
"Plan calls to comparison and metrics tools",
"Emit JSON matching AnalysisOutputSchema"
],
output_instructions=["Return JSON exactly matching AnalysisOutputSchema"]
)
analysis_agent = BaseAgent(
config=BaseAgentConfig(
client=client,
model=model,
system_prompt_generator=analysis_system_prompt,
memory=AgentMemory(),
input_schema=AnalysisInputSchema,
output_schema=AnalysisOutputSchema
)
)
Now, regarding the orchestrator, I am not sure how to setup! ( i have checked the example but I can't figure)
So, I can't figure how to
- setup the orchstrator in order to decide which agent it should call
- make each agent to have access to the specific tool/tools
- can I just use the python funtions as tools or do i have ot create a tool class?
Thank you!
On first sight your agents look good, what your example here is missing is the overarching orchestrator agent, but instead of the tool input schemas, you would be doing a Union of the agent input schemas...
In hindsight, the example could/should be more clear and straightforward and show both tool-calling & agent-calling (which this framework essentially treats as the same thing)
I'll work on that, in the meanwhile feel free to join the discord server if you need help with anything else!
make each agent to have access to the specific tool/tools
Basically, whenever the output schema of an agent is:
- The input schema of a tool
- A
Unionof multiple Tool input schemas
Let's pretend we have an agent my_agent and you have defined its input schema like this:
class MyAgentInputSchema(BaseIOSchema):
'''The input schema of the agent'''
user_message: str = Field(..., description="The input of the user")
If we only want to use the calculator tool, for example with this agent, this agent does not need an output schema, since we want the agent to output the input schema of the calculator tool...
You can just do something like:
my_agent = BaseAgent(
config=BaseAgentConfig(
...
input_schema=MyAgentInputSchema,
output_schema=CalculatorInputSchema
)
)
calculator_tool_input = my_agent.run("Please sum up the product of 123 and 321 together")
calculator_tool_output = calculator_tool.run(calculator_tool_input)
So, following this pattern, it doesn't matter if the output schema of my_agent is a tool or an agent, it's all based on input & output schemas and a run() method
can I just use the python funtions as tools or do i have ot create a tool class?
So, for the above pattern to work, every tool should be a tool class, all you need is to define the input & output schema, and a run() method
Thank you Kenny for the answer!
Just to be clear since I am reading what you said again.
You say that if for example the my_agent wants only to use the calculator tool , there is no need for us to create an output schema for my_agent because it wants to output the input of calculator tool.
So, we do not define an MyAgentOutputSchema for my_agent.
But, when we create the my_agent with :
my_agent = BaseAgent(
config=BaseAgentConfig(
...
input_schema=MyAgentInputSchema,
output_schema=CalculatorInputSchema
)
)
we do use the output_schema there by placing the CalculatorInputSchema.
Finally, we don't use the CalculatorOutputSchema somewhere (note I mention the Output), but
we still receive the results from the calculator tool (which invokes CalculatorOutputSchema) because we create:
calculator = CalculatorTool() , which CalculatorTool calls both input and output schemas.
right?
Is the above correct?
And finally, if an agent wants to use 2 or more tools?
Just place for example output_schema=CalculatorInputSchema, AnotherToolInputSchema ?
Or we need to create something like:
def execute_tool(
searxng_tool: SearxNGSearchTool, calculator_tool: CalculatorTool, orchestrator_output: OrchestratorOutputSchema
) -> Union[SearxNGSearchToolOutputSchema, CalculatorToolOutputSchema]:
if orchestrator_output.tool == "search":
return searxng_tool.run(orchestrator_output.tool_parameters)
elif orchestrator_output.tool == "calculator":
return calculator_tool.run(orchestrator_output.tool_parameters)
else:
raise ValueError(f"Unknown tool: {orchestrator_output.tool}")
like you have in the orchestrator examples?
So kind of creating an "orhestrator" just for the 2 tools?
we do use the output_schema there by placing the CalculatorInputSchema.
Yes, absolutely correct, this is also demonstrated for example in the query generation agent in this example
Here you see that the desired output must be the input format of the search tool. If at any point we decide to use another search tool that requires a slightly different input, say one tool's input schema has a category field like general, news, social and the other tool does not, then you don't really need to change any code, purely based on the input & output schemas, the agent will know that it now needs to include the category news when I ask "Give me the latest updates on ...." for tool one, but not to do so if we switch out the tool... Does that make sense?
Is the above correct?
Yes, each agent & tool has an input & output schema and you are correct in its usage!
This way you can also chain agents & tools together, this may be a bit of a silly example but, imagine you have:
- An agent whose only job is to take a user's input, like "find 5-star-rated sushi restaurants in Brussels" & convert it to a yelp API call
- A yelp API tool
- An agent whose only job is to take a yelp API call's output and give results to the user in a chat message
Then you would have: AgentOne with ChatMessage as its INPUT & YelpToolInputSchema as its OUTPUT-> YelpTool with YelpToolInputSchema as its INPUT & YelpToolOutputSchema as OUTPUT -> AgentTwo with YelpToolOutputSchema as its INPUT and finally a ChatMessage output
And finally, if an agent wants to use 2 or more tools? Just place for example output_schema=CalculatorInputSchema, AnotherToolInputSchema ?
What you'd want to do here is use the Union type
class OrchestratorOutputSchema(BaseIOSchema):
"""Combined output schema for the Orchestrator Agent. Contains the tool to use and its parameters."""
tool: str = Field(..., description="The tool to use: 'search' or 'calculator'")
tool_parameters: Union[SearxNGSearchToolInputSchema, CalculatorToolInputSchema] = Field(
..., description="The parameters for the selected tool"
)
Now, it has been a while so I am a tiny bit vague on this example, but I am fairly certain I only included tool: str = Field(..., description="The tool to use: 'search' or 'calculator'") because later on, it would be a bit easier to check the tool name in the if/else statement and also the success rate might have been a bit higher with weaker models (I believe this example already exists since gpt-4-turbo)
But in essence, it can simply be:
class OrchestratorOutputSchema(BaseIOSchema):
"""Combined output schema for the Orchestrator Agent. Contains the tool to use."""
tool: Union[SearxNGSearchToolInputSchema, CalculatorToolInputSchema] = Field(
..., description="The parameters for the selected tool"
)
You could then have some code that, instead of doing an if/else on the tool name, just checks the type of the tool property, which contains the parameters of the tool to call directly... (Look at the Union as a list of possible output types, once you have ran the agent, it can only be either the SearxNGSearchToolInputSchema or the CalculatorToolInputSchema
And of course, in Atomic Agents, agents and tools are very interchangeable, so you can even do:
class OrchestratorOutputSchema(BaseIOSchema):
"""Combined output schema that decides whether to use a tool or call another agent."""
tool: Union[SearxNGSearchToolInputSchema, CalculatorToolInputSchema, SomeOtherAgent, AnotherOtherAgent] = Field(
..., description="The parameters for the selected tool or agent"
)
Another cool thing you can do, which is done in the MCP example is have a Union between tool schemas and a "Final output schema", which allows you to make agents more autonomous, but that's a little off-topic here though I guess I should make a more clear non-MCP example for that