Routing

There are two ways to implement routing in AgentScope, both simple and easy to implement:

  • Routing by structured output

  • Routing by tool calls

Tip

Considering there is no unified standard/definition for agent routing, we follow the setting in Building effective agents

Routing by Structured Output

By this way, we can directly use the structured output of the agent to determine which agent to route the message to.

Initialize a routing agent

import asyncio
import json
import os
from typing import Literal

from pydantic import BaseModel, Field

from agentscope.agent import ReActAgent
from agentscope.formatter import DashScopeChatFormatter
from agentscope.memory import InMemoryMemory
from agentscope.message import Msg
from agentscope.model import DashScopeChatModel
from agentscope.tool import Toolkit, ToolResponse

router = ReActAgent(
    name="Router",
    sys_prompt="You're a routing agent. Your target is to route the user query to the right follow-up task.",
    model=DashScopeChatModel(
        model_name="qwen-max",
        api_key=os.environ["DASHSCOPE_API_KEY"],
        stream=False,
    ),
    formatter=DashScopeChatFormatter(),
)


# Use structured output to specify the routing task
class RoutingChoice(BaseModel):
    your_choice: Literal[
        "Content Generation",
        "Programming",
        "Information Retrieval",
        None,
    ] = Field(
        description="Choose the right follow-up task, and choose ``None`` if the task is too simple or no suitable task",
    )
    task_description: str | None = Field(
        description="The task description",
        default=None,
    )


async def example_router_explicit() -> None:
    """Example of explicit routing with structured output."""
    msg_user = Msg(
        "user",
        "Help me to write a poem",
        "user",
    )

    # Route the query
    msg_res = await router(
        msg_user,
        structured_model=RoutingChoice,
    )

    # The structured output is stored in the metadata field
    print("The structured output:")
    print(json.dumps(msg_res.metadata, indent=4, ensure_ascii=False))


asyncio.run(example_router_explicit())
Router: {
    "type": "tool_use",
    "name": "generate_response",
    "input": {
        "your_choice": "Content Generation",
        "task_description": "Help me to write a poem"
    },
    "id": "call_ab3d69680a5e4075882e33"
}
system: {
    "type": "tool_result",
    "id": "call_ab3d69680a5e4075882e33",
    "name": "generate_response",
    "output": [
        {
            "type": "text",
            "text": "Successfully generated response."
        }
    ]
}
Router: Sure, I can help with that! Could you please provide more details? For example, do you have a specific theme or style in mind for the poem? Any particular emotions or messages you'd like to convey? The more information you give me, the better I can tailor the poem to your liking.
The structured output:
{
    "your_choice": "Content Generation",
    "task_description": "Help me to write a poem"
}

Routing by Tool Calls

Another way is to wrap the downstream agents into a tool function, so that the routing agent decides which tool to call based on the user query.

We first define several tool functions:

async def generate_python(demand: str) -> ToolResponse:
    """Generate Python code based on the demand.

    Args:
        demand (``str``):
            The demand for the Python code.
    """
    # An example demand agent
    python_agent = ReActAgent(
        name="PythonAgent",
        sys_prompt="You're a Python expert, your target is to generate Python code based on the demand.",
        model=DashScopeChatModel(
            model_name="qwen-max",
            api_key=os.environ["DASHSCOPE_API_KEY"],
            stream=False,
        ),
        memory=InMemoryMemory(),
        formatter=DashScopeChatFormatter(),
        toolkit=Toolkit(),
    )
    msg_res = await python_agent(Msg("user", demand, "user"))

    return ToolResponse(
        content=msg_res.get_content_blocks("text"),
    )


# Fake some other tool functions for demonstration purposes
async def generate_poem(demand: str) -> ToolResponse:
    """Generate a poem based on the demand.

    Args:
        demand (``str``):
            The demand for the poem.
    """
    pass


async def web_search(query: str) -> ToolResponse:
    """Search the web for the query.

    Args:
        query (``str``):
            The query to search.
    """
    pass

After that, we define a routing agent and equip it with the above tool functions.

toolkit = Toolkit()
toolkit.register_tool_function(generate_python)
toolkit.register_tool_function(generate_poem)
toolkit.register_tool_function(web_search)

# Initialize the routing agent with the toolkit
router_implicit = ReActAgent(
    name="Router",
    sys_prompt="You're a routing agent. Your target is to route the user query to the right follow-up task.",
    model=DashScopeChatModel(
        model_name="qwen-max",
        api_key=os.environ["DASHSCOPE_API_KEY"],
        stream=False,
    ),
    formatter=DashScopeChatFormatter(),
    toolkit=toolkit,
    memory=InMemoryMemory(),
)


async def example_router_implicit() -> None:
    """Example of implicit routing with tool calls."""
    msg_user = Msg(
        "user",
        "Help me to generate a quick sort function in Python",
        "user",
    )

    # Route the query
    await router_implicit(msg_user)


asyncio.run(example_router_implicit())
Router: {
    "type": "tool_use",
    "name": "generate_python",
    "input": {
        "demand": "a quick sort function"
    },
    "id": "call_84bdc23ca64a435bb3282a"
}
PythonAgent: Certainly! Below is a Python implementation of the Quick Sort algorithm. This implementation uses the Lomuto partition scheme, which is a common and straightforward approach.

```python
def quick_sort(arr, low, high):
    if low < high:
        # Partition the array
        pi = partition(arr, low, high)

        # Recursively sort elements before and after partition
        quick_sort(arr, low, pi - 1)
        quick_sort(arr, pi + 1, high)

def partition(arr, low, high):
    pivot = arr[high]  # Choose the last element as the pivot
    i = low - 1  # Pointer for the greater element

    for j in range(low, high):
        # If current element is smaller than or equal to pivot
        if arr[j] <= pivot:
            i += 1  # Increment the pointer
            # Swap elements at i and j
            arr[i], arr[j] = arr[j], arr[i]

    # Swap the pivot element with the element at i+1
    arr[i + 1], arr[high] = arr[high], arr[i + 1]
    return i + 1

# Example usage
if __name__ == "__main__":
    arr = [10, 7, 8, 9, 1, 5]
    n = len(arr)
    quick_sort(arr, 0, n - 1)
    print("Sorted array:", arr)
```

### Explanation:
1. **quick_sort function**:
   - Takes an array `arr`, and two indices `low` and `high`.
   - If `low` is less than `high`, it partitions the array and recursively sorts the sub-arrays.

2. **partition function**:
   - Takes an array `arr`, and two indices `low` and `high`.
   - Chooses the last element as the pivot.
   - Reorders the array such that all elements less than or equal to the pivot are on its left, and all elements greater than the pivot are on its right.
   - Returns the index of the pivot after partitioning.

3. **Example usage**:
   - An example array is defined and sorted using the `quick_sort` function.
   - The sorted array is printed.

You can run this code to see the Quick Sort algorithm in action.
system: {
    "type": "tool_result",
    "id": "call_84bdc23ca64a435bb3282a",
    "name": "generate_python",
    "output": [
        {
            "type": "text",
            "text": "Certainly! Below is a Python implementation of the Quick Sort algorithm. This implementation uses the Lomuto partition scheme, which is a common and straightforward approach.\n\n```python\ndef quick_sort(arr, low, high):\n    if low < high:\n        # Partition the array\n        pi = partition(arr, low, high)\n        \n        # Recursively sort elements before and after partition\n        quick_sort(arr, low, pi - 1)\n        quick_sort(arr, pi + 1, high)\n\ndef partition(arr, low, high):\n    pivot = arr[high]  # Choose the last element as the pivot\n    i = low - 1  # Pointer for the greater element\n\n    for j in range(low, high):\n        # If current element is smaller than or equal to pivot\n        if arr[j] <= pivot:\n            i += 1  # Increment the pointer\n            # Swap elements at i and j\n            arr[i], arr[j] = arr[j], arr[i]\n\n    # Swap the pivot element with the element at i+1\n    arr[i + 1], arr[high] = arr[high], arr[i + 1]\n    return i + 1\n\n# Example usage\nif __name__ == \"__main__\":\n    arr = [10, 7, 8, 9, 1, 5]\n    n = len(arr)\n    quick_sort(arr, 0, n - 1)\n    print(\"Sorted array:\", arr)\n```\n\n### Explanation:\n1. **quick_sort function**:\n   - Takes an array `arr`, and two indices `low` and `high`.\n   - If `low` is less than `high`, it partitions the array and recursively sorts the sub-arrays.\n\n2. **partition function**:\n   - Takes an array `arr`, and two indices `low` and `high`.\n   - Chooses the last element as the pivot.\n   - Reorders the array such that all elements less than or equal to the pivot are on its left, and all elements greater than the pivot are on its right.\n   - Returns the index of the pivot after partitioning.\n\n3. **Example usage**:\n   - An example array is defined and sorted using the `quick_sort` function.\n   - The sorted array is printed.\n\nYou can run this code to see the Quick Sort algorithm in action."
        }
    ]
}
Router: The Python code for a quick sort function has been generated. Here it is:

```python
def quick_sort(arr, low, high):
    if low < high:
        # Partition the array
        pi = partition(arr, low, high)

        # Recursively sort elements before and after partition
        quick_sort(arr, low, pi - 1)
        quick_sort(arr, pi + 1, high)

def partition(arr, low, high):
    pivot = arr[high]  # Choose the last element as the pivot
    i = low - 1  # Pointer for the greater element

    for j in range(low, high):
        # If current element is smaller than or equal to pivot
        if arr[j] <= pivot:
            i += 1  # Increment the pointer
            # Swap elements at i and j
            arr[i], arr[j] = arr[j], arr[i]

    # Swap the pivot element with the element at i+1
    arr[i + 1], arr[high] = arr[high], arr[i + 1]
    return i + 1

# Example usage
if __name__ == "__main__":
    arr = [10, 7, 8, 9, 1, 5]
    n = len(arr)
    quick_sort(arr, 0, n - 1)
    print("Sorted array:", arr)
```

This implementation of Quick Sort uses the Lomuto partition scheme, where the last element is chosen as the pivot. The `partition` function reorders the elements so that all elements less than or equal to the pivot are on its left, and all elements greater than the pivot are on its right. The `quick_sort` function then recursively sorts the sub-arrays.

You can run this code to see the Quick Sort algorithm in action.

Total running time of the script: (0 minutes 36.350 seconds)

Gallery generated by Sphinx-Gallery