Customize Orchestrator for Task Assignments¶

Overview¶

The Orchestrator is a general-purpose component that routes users' requests i.e., queries to the appropriate agents. It ensures that tasks are assigned to the correct agents based on their capabilities and the nature of the request. Additionally, it allows features with decomposition and customizable prompt and contexts.

Goals¶

This tutorial will guide you through the following steps:

Get an overview of Orchestrator and its workflow.
Create or modify a YAML configuration file.
Show examples of Orchestrator with customized configurations.

Orchestrator Workflow¶

RAI Check (Optional) – Every query first passes through a Responsible-AI compliance check; unsafe queries are rejected.
Task Plan Generation – If enabled, the query is broken down into subtasks with relevant context; otherwise it is routed directly.
Task Execution – The Orchestrator assigns tasks to the appropriate agents and executes them.

Configuration¶

To customize Orchestrator, you need to define it in the YAML configuration. This configuration specifies the Orchestrator behavior.

Configuration Parameters¶

agent_list (required): List of agents the Orchestrator can route tasks to.
- Each entry must include agent_name, which must match a name defined in the utility_agents section.
enable_routing (optional, default: true): Controls LLM-based routing.
- true: Queries may be decomposed and routed to different agents.
- false: Every query is sent directly to the first agent in agent_list.
decompose (optional, default: true): Controls whether queries are decomposed into subtasks.
- true: The Orchestrator may split queries into subtasks.
- false: Queries are routed as-is without decomposition.
rai_config (optional): Override configuration for the Responsible-AI (RAI) engine.
system_prompt_suffix (optional): Extra instructions appended to the system prompt to guide Orchestrator behavior. This field allows developers to add application-specific guidance or constraints.
contexts (optional): Defines which contexts are included in the Orchestrator prompt. Options:
- "date" – Inserts today’s date.
- "env_variable" – Includes key–value pairs for agent context, which can be static or updated at runtime.
- "chat_history" – Includes recent conversation turns.
- "relevant_chat_history" – Includes the most relevant chat history retrieved by semantic search.

Here’s an example configuration:

orchestrator:
  agent_list:                  # Required. List of agents the orchestrator can route to.
    - agent_name: "Search Agent"   # Must match a utility agent defined elsewhere.
    - agent_name: "Planner Agent" # Multiple agents can be listed.
  system_prompt_suffix: "Ensure the dinner cost stays under the value of dinner_budget."  
                               # Optional. Extra instruction appended to the system prompt.
  contexts:                    # Optional. Context sources included when building the prompt.
    - "date"                   # Inserts today’s date.
    - "env_variable"           # Includes key–value pairs for agent context, which can be static or updated at runtime.
    - "chat_history"           # Adds recent conversation turns.
    - "relevant_chat_history"  # Retrieves the most relevant past chat history.

Example Usage¶

This section demonstrates how to customize the Orchestrator. Details on customizing rai_config are available in Responsible AI Module.

1. YAML Configuration Files¶

To enable a customized Orchestrator, you need to define it in a YAML file.

a. Routing and Decomposition¶

By default, the Orchestrator enables both routing and decomposition.

When enable_routing is true and decompose is true, the query will be split into sub-queries.
When enable_routing is true and decompose is false, the original query is routed.

When enable_routing is false, the query is sent directly to the first agent in agent_list

orchestrator:
  agent_list:                                    # Required. List of agents orchestrator can route tasks to
    - agent_name: "Planner Agent"                # Must match a utility agent name below
    - agent_name: "Search Agent"                 # Multiple agents can be listed
  enable_routing: true
  decompose: true

utility_agents:
  - agent_class: PlanningAgent                   # Define Planner Agent
    agent_name: "Planner Agent"                  # Name used by orchestrator
    agent_description: "Turns constraints and facts into a concrete, time-ordered plan with rationale and trade-offs."

  - agent_class: SearchAgent                     # Define Search Agent
    agent_name: "Search Agent"                   # Name used by orchestrator
    agent_description: "Search for reliable facts (hours, prices, locations, travel times)."

To disable the routing, explicitly set enable_routing to false. The query will be sent directly to the first agent in agent_list. In the example below, it will be sent to Planner Agent.

orchestrator:
  agent_list:                                    # Required. List of agents orchestrator can route tasks to
    - agent_name: "Planner Agent"                # Must match a utility agent name below
    - agent_name: "Search Agent"                 # Multiple agents can be listed
  enable_routing: false

utility_agents:
  - agent_class: PlanningAgent                   # Define Planner Agent
    agent_name: "Planner Agent"                  # Name used by orchestrator
    agent_description: "Turns constraints and facts into a concrete, time-ordered plan with rationale and trade-offs."

  - agent_class: SearchAgent                     # Define Search Agent
    agent_name: "Search Agent"                   # Name used by orchestrator
    agent_description: "Search for reliable facts (hours, prices, locations, travel times)."

b. Adding `system_prompt_suffix`¶

The Orchestrator provides general task-routing logic. You can append task-specific guidance using system_prompt_suffix, which is added to the base system prompt.

orchestrator:
  agent_list:                                    # Required. List of agents orchestrator can route tasks to
    - agent_name: "Planner Agent"                # Must match a utility agent name below
    - agent_name: "Search Agent"                 # Multiple agents can be listed
  system_prompt_suffix: "Encourage Planner Agent and Search Agent to work together. Planner should propose the itinerary structure, Search agent should provide facts to support it, and Planner should refine using those facts. A suggested workflow would be Planner Agent, Search Agent, Planner Agent."

utility_agents:
  - agent_class: PlanningAgent                   # Define Planner Agent
    agent_name: "Planner Agent"                  # Name used by orchestrator
    agent_description: "Turns constraints and facts into a concrete, time-ordered plan with rationale and trade-offs."

  - agent_class: SearchAgent                     # Define Search Agent
    agent_name: "Search Agent"                   # Name used by orchestrator
    agent_description: "Search for reliable facts (hours, prices, locations, travel times)."

c. Using `contexts`¶

This configuration example supports the following scenario:

The Orchestrator coordinates between two agents:
- SearchAgent – retrieves reliable facts (e.g., hours, prices, locations, travel times).
- PlannerAgent – organizes those facts into a concrete, time-ordered plan with rationale and trade-offs.
system_prompt_suffix – ensures plans respect a specified budget.
contexts – provide background information such as environment variables and chat history.

memory_config:  
  memory_modules:  
    - memory_name: env_variable                  # Define a memory module for environment variables
      memory_class: VariableMemoryModule
      kwargs:  
        variables:                               # Store custom variables for use in prompts
          dinner_budget: "$100 per person"       # Example variable: dinner budget constraint

orchestrator:
  agent_list:                                    # Required. List of agents orchestrator can route tasks to
    - agent_name: "Planner Agent"                # Must match a utility agent name below
    - agent_name: "Search Agent"                 # Multiple agents can be listed
  system_prompt_suffix: "Ensure the dinner cost stays under the value of dinner_budget."
                                                 # Optional. Extra instruction appended to system prompt
  contexts:                                      # Optional. Context sources included in orchestrator prompt
    - "date"                                     # Inserts today’s date
    - "env_variable"                             # Includes environment variables (e.g., dinner_budget)
    - "chat_history"                             # Adds recent conversation turns
    - "relevant_chat_history"                    # Retrieves the most relevant past chat history

utility_agents:
  - agent_class: PlanningAgent                   # Define Planner Agent
    agent_name: "Planner Agent"                  # Name used by orchestrator
    agent_description: "Turns constraints and facts into a concrete, time-ordered plan with rationale and trade-offs."

  - agent_class: SearchAgent                     # Define Search Agent
    agent_name: "Search Agent"                   # Name used by orchestrator
    agent_description: "Search for reliable facts (hours, prices, locations, travel times)."

2. Python File¶

Now, you can start the development using these lines of code:

Python Code¶

import asyncio
import os

from air import DistillerClient
from air.utils import async_print
from dotenv import load_dotenv

load_dotenv() # loads your API_KEY from your local '.env' file
api_key=str(os.getenv("API_KEY"))

async def main():
    """
    Runs the customizable orchestrator demo.
    """
    # Initialize Distiller client, project name, and session id
    client = DistillerClient(api_key=api_key)
    project_name = "orchestrator_project"
    session_uuid = f"session_{os.getpid()}"

    # Initialize the orchestrator project
    client.create_project(config_path="config.yaml", project=project_name)

    async with client(project=project_name, uuid=session_uuid) as dc:
        query = "Plan a 1-day itinerary in Seattle for Saturday. I want to visit an art museum in the morning, a famous viewpoint around sunset, and have a seafood dinner. Please ensure opening hours and rough costs are correct."
        responses = await dc.query(query=query)

        print(f"--- Running Query: {query} ---")
        async for response in responses:
            await async_print(
                f"Response from {response['role']}: {response['content']}"
            )

        # Clear session memory after the run
        await dc.reset_memory()
        await async_print("--- Session Complete ---")


if __name__ == "__main__":
    asyncio.run(main())

(Optional) Python Code for updating the env_variable at runtime¶

In addition to being initialized in the YAML file, the env_variable can also be initialized or updated at runtime. Following shows an example.

import asyncio
import os

from air import DistillerClient
from air.utils import async_print
from dotenv import load_dotenv

load_dotenv() # loads your API_KEY from your local '.env' file
api_key=str(os.getenv("API_KEY"))

async def main():
    """
    Runs the customizable orchestrator demo.
    """
    # Initialize Distiller client, project name, and session id
    client = DistillerClient(api_key=api_key)
    project_name = "orchestrator_project"
    session_uuid = f"session_{os.getpid()}"

    # Initialize the orchestrator project
    client.create_project(config_path="config.yaml", project=project_name)

    async with client(project=project_name, uuid=session_uuid) as dc:
        # Add runtime environment variables to memory
        await dc.add_memory(
            source="env_variable",
            variables_dict={
                "dinner_budget": "$50 per person",
            },
        )

        query = "Plan a 1-day itinerary in Seattle for Saturday. I want to visit an art museum in the morning, a famous viewpoint around sunset, and have a seafood dinner. Please ensure opening hours and rough costs are correct."
        responses = await dc.query(query=query)

        print(f"--- Running Query: {query} ---")
        async for response in responses:
            await async_print(
                f"Response from {response['role']}: {response['content']}"
            )

        # Clear session memory after the run
        await dc.reset_memory()
        await async_print("--- Session Complete ---")


if __name__ == "__main__":
    asyncio.run(main())

Sample Outputs¶

--- Running Query: Plan a 1-day itinerary in Seattle for Saturday. I want to visit an art museum in the morning, a famous viewpoint around sunset, and have a seafood dinner. Please ensure opening hours and rough costs are correct. ---
Response from orchestrator: I will decompose your query into subtasks and handle them one by one. Search Agent, opening hours and rough costs of art museums in Seattle
Response from Search Agent: Searching over Web Search
Response from Search Agent: I've got the information you need about art museums in Seattle. The Seattle Art Museum and the Seattle Asian Art Museum are two popular options. [Output abbreviated]

Response from orchestrator: Search Agent, famous viewpoints in Seattle for sunset
Response from Search Agent: Searching over Web Search
Response from Search Agent: For a famous viewpoint around sunset, I recommend visiting the Space Needle Observation Deck. [Output abbreviated]

Response from orchestrator: Search Agent, seafood restaurants in Seattle with dinner cost under $100 per person
Response from Search Agent: Searching over Web Search
Response from Search Agent: **Seafood Restaurants in Seattle with Dinner Cost under $100 per Person** [Output abbreviated]

Response from orchestrator: Planner Agent, Plan a 1-day itinerary in Seattle for Saturday. Visit an art museum in the morning, a famous viewpoint around sunset, and have a seafood dinner. Consider the results from previous searches.
Response from Planner Agent: I'd be happy to help you plan a day in Seattle. Based on previous searches, I recommend starting your day at the Frye Art Museum, which is open from 11 am to 5 pm on Saturdays. [Output abbreviated]
--- Session Complete ---