Feature: MapAgent (#1828)

[ananyablonko]

• Implements: New Workflow Agent: MapAgent - calls 1 subagent asynchronously multiple times with different arguments #1828

Adds basic functionality of MapAgent.

ADK misses a map functionality, where a single agent operates in parallel on multiple data.
For example, if I want to create N images in parallel, where N is only known at runtime, this cannot be done using ADK.

To circumvent this, one might use a Parallel Agent with N agents that are clones of one another, and create the agent dynamically. This causes problems as explained in the linked feature request.

A MapAgent comes to the rescue!

Sample Usage of MapAgent:

import asyncio
from pydantic import RootModel
from google.adk.agents import LlmAgent, SequentialAgent, MapAgent
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import types
import logging 


logging.basicConfig( filename='log.log', level=logging.DEBUG)

title_agent = LlmAgent(
    name="title_agent",
    model="gemini-2.0-flash",
    output_schema=RootModel[list[str]]
)

story_agent = LlmAgent(
    name="story_agent",
    model="gemini-2.0-flash",
    instruction="Write a short story for the given title. Output the story and nothing else.",
    include_contents="none",
)

map_agent = MapAgent(
    name="map_agent",
    sub_agents=[story_agent]
)

root = SequentialAgent(
    name="root",
    sub_agents=[title_agent, map_agent]
)

runner = Runner(
    agent=root,
    app_name="MapAgentTest",
    session_service=InMemorySessionService(),
)

async def main():
    message = types.Content(
        role="user",
        parts=[types.Part(text="Give me three sci-fi story titles")]
    )
    await runner.session_service.create_session(app_name=runner.app_name, user_id="0", session_id="0")

    async for event in runner.run_async(
        user_id="0",
        session_id="0",
        new_message=message,
    ):
        if event.is_final_response():
            print(((event.content or types.Content).parts or [types.Part()])[0].text)

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

The PR includes 100% coverage unittests that use all kinds of agents as sub agents or parents of the MapAgent.

[gemini-code-assist]

Summary of Changes

Hello @ananyablonko, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces the MapAgent feature to the ADK framework, addressing the current limitation of processing multiple data inputs in parallel with a single agent instance. The MapAgent streamlines the creation of dynamic, parallel workflows by allowing a single sub-agent to be cloned and executed for each item in a list of prompts, ensuring efficient and scalable agent operations within complex agent compositions.

Highlights

• Introduction of MapAgent: A new MapAgent class has been added to the ADK framework, designed to enable a single agent to operate in parallel on multiple data points, addressing limitations in handling runtime-known N parallel operations.
• Dynamic Sub-Agent Management: The MapAgent dynamically creates and manages unique copies of a specified sub-agent for each input prompt, ensuring isolated and parallel execution. It handles branching of invocation contexts and agent trees to maintain thread safety.
• Input Processing and Event Handling: The MapAgent extracts a list of prompts (serialized as a JSON string) from the latest session event. It temporarily removes this event for sub-agent invocations to prevent interference between sibling agents and adds individual prompts as temporary events for each branched context.
• Comprehensive Unit Testing: Extensive unit tests have been added, achieving 100% coverage for the MapAgent. These tests validate its functionality in various configurations, including empty input, basic text input, and complex agent hierarchies involving LoopAgent, ParallelAgent, SequentialAgent, and nested MapAgents.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces the MapAgent, a new agent that applies a single sub-agent to multiple inputs in parallel. This is a great addition to the ADK, addressing a need for parallel processing of dynamic data sets. The implementation is clean and leverages existing components like _merge_agent_run from ParallelAgent. The unit tests are comprehensive, covering various agent compositions and edge cases.

I've added a couple of comments to improve code style and maintainability. One is about using proper comments in the agent implementation, and the other suggests refactoring duplicated code in the tests by using an existing helper function. Overall, this is a solid contribution.

[ananyablonko]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces the MapAgent, a new agent that applies a sub-agent to a list of inputs in parallel. The implementation is solid and the accompanying unit tests provide good coverage for various agent compositions. I've identified a critical bug in how sub-agents are cloned, which could lead to issues with parent_agent references. I've also suggested an improvement to make input parsing more robust against malformed data. Overall, this is a great addition to the ADK.

[gemini-code-assist]

The current implementation assumes the input message is always a valid JSON-serialized list of strings. If input_message is an empty string or not a valid JSON, model_validate_json will raise a ValidationError, causing the agent to crash. It's better to handle this potential error gracefully.

I suggest wrapping this call in a try...except block to catch ValidationError. If parsing fails, you could log a warning and default to an empty list of prompts, making the agent more robust.

You'll also need to add import logging, from pydantic import ValidationError, and initialize a logger: logger = logging.getLogger("google_adk." + __name__) at the top of the file.

Suggested change

	agent_input = RootModel[list[str]].model_validate_json(input_message).root
	try:
	agent_input = RootModel[list[str]].model_validate_json(input_message).root
	except ValidationError:
	logger.warning("Invalid input for MapAgent, expected a JSON list of strings, got: %s", input_message)
	agent_input = []

[ananyablonko]

Respectfully, I believe invoking this agent with a non-list input is a developer error. In that case, wrong input should cause a crash.
For now, I added a test to ensure this behavior.

[ryanaiagent]

Hi @ananyablonko , Thank you for your contribution! We appreciate you taking the time to submit this pull request.
Can you please fix the failing unit tests before we can proceed with the review.