DEV Community: Henry Ruiz, PhD

Post-Training Gemma 3 for Earth Observation (EO) Understanding: A JAX Stack + TPU Pipeline for…

Henry Ruiz, PhD — Mon, 20 Apr 2026 17:12:30 +0000

Post-Training Gemma 3 for Earth Observation (EO) Understanding: A JAX Stack + TPU Pipeline for Multi-Label Sentinel Satellite Remote Sensing Scene Classification

Earth observation has no shortage of data. What it still needs is broader, more practical adoption of scalable approaches to adapt modern multimodal models to geospatial tasks — without turning training into a months-long infrastructure project.

That was the motivation behind building GemmaEarth as part of the TPU Sprint.

GitHub - haruiz/gemmaearth

GemmaEarth is a domain-focused post-training and benchmarking pipeline built with Tunix, designed to adapt Google’s Gemma 3 4B IT model for Earth Observation (EO) understanding using the EarthDial: Turning Multi-sensory Earth Observations into Interactive Dialogues (2025) dataset (https://arxiv.org/abs/2412.15190). EarthDial provides a large-scale instruction-tuning dataset spanning diverse EO tasks, including classification, captioning, visual question answering (VQA), and reasoning.

While the long-term goal is to fine-tune the model across the full range of EarthDial tasks, TPU resource allocation constraints led this project to begin with a focused and practical use case: satellite remote-sensing scene classification. This initial stage leverages the BigEarthNet subset, a large-scale Sentinel benchmark for multi-label land-use and land-cover classification.

EarthDial dataset

Project Contributions

GemmaEarth provides a reproducible, end-to-end workflow for adapting Gemma 3 4B IT to Earth Observation tasks, with a primary focus on multi-label satellite scene classification. The pipeline is designed to be extensible to a broader set of EO tasks within the EarthDial dataset, leveraging a TPU-native JAX stack.

Rather than focusing on a single training run, the project emphasizes a reference workflow spanning data preparation, post-training, evaluation, and benchmarking.

The implementation leverages the full JAX ecosystem (https://jaxstack.ai/):

Tunix orchestrates the PEFT-based training pipeline
Grain enables scalable data processing
Optax handles optimization and scheduling
Orbax manages checkpointing
Qwix injects LoRA adapters for efficient fine-tuning

All experiments were executed on Google Cloud TPU v5litepod-8, with an emphasis on scalability and reproducibility.

Key contributions include:

A reproducible JAX-based pipeline for end-to-end post-training and benchmarking
Parameter-efficient adaptation of Gemma 3 using LoRA in a TPU-sharded setup
Multimodal training support using image–text EO data
TPU-native data and training pipelines for scalable experimentation
A practical evaluation framework with interpretable metrics for multi-label classification

JAX AI stack

The JAX Stack in Practice

GemmaEarth integrates multiple components from the JAX ecosystem into a cohesive, TPU-native training pipeline, structured across key layers: data processing, model adaptation, distributed training, optimization, and checkpointing.

Data Loading and Preprocessing

Efficient data handling is critical for TPU utilization. In GemmaEarth, Grain is used to construct a high-performance, composable input pipeline for multimodal data.

This pipeline handles shuffling, transformation, batching, and iteration, ensuring a steady stream of preprocessed data into the training loop.

Source: https://github.com/haruiz/gemma_earth/blob/5fd2e4c0b6cf3c65eac026bd5aecc136e81072b8/src/gemma_earth/dataset.py#L519

def _build_train_pipeline(
        self,
        split: Any,
        image_processor: Any,
        tokenizer: tokenizer_lib.Tokenizer,
        num_epochs: int | None,
        split_name: str,
    ) -> Any:
        """Construct a shuffled, mapped, and batched Grain iterable pipeline.

        The pipeline applies the following operations in order: shuffle, map to
        training example, map to training input (tokenize + pad), batch, repeat,
        and convert to an iterable dataset.

        Args:
            split: HuggingFace Dataset split used as the Grain data source.
            image_processor: Tunix ImageProcessor forwarded to
                _to_training_example.
            tokenizer: Tunix tokenizer forwarded to _to_training_input.
            num_epochs: Number of times to repeat the dataset. Pass 1 (or
                None) for a single pass, typically used for validation.
            split_name: Human-readable label ("train" or "validation")
                used in progress log messages.

        Returns:
            A Grain iterable dataset ready for use as train_ds or
            eval_ds in GemmaEarth.train.
        """
        settings = self.settings
        logger.info("Building %s pipeline with %d rows...", split_name, len(split))
        return (
            grain.MapDataset.source(split)
            .shuffle(seed=settings.shuffle_seed)
            .map(lambda x: self._to_training_example(x, image_processor))
            .map(lambda x: self._to_training_input(x, tokenizer))
            .batch(settings.batch_size, drop_remainder=True)
            .repeat(num_epochs)
            .to_iter_dataset()
        )

Distributed Training (Tunix + Jax)

In GemmaEarth, Tunix orchestrates the training loop, coordinating model updates, data flow, and optimization across the entire stack. It provides a unified interface for:

State-of-the-art TPU performance through JAX-native execution
Supervised Fine-Tuning (SFT) for instruction-based adaptation
Reinforcement Learning (RL) for policy optimization
Agentic RL , enabling training loops where agents interact with dynamic environments

Source: https://github.com/haruiz/gemma_earth/blob/5fd2e4c0b6cf3c65eac026bd5aecc136e81072b8/src/gemma_earth/trainers/base.py#L386

def _build_trainer(
        self,
        lora_model: nnx.Module,
        optimizer: optax.GradientTransformation,
        max_steps: int,
        tokenizer: tokenizer_lib.Tokenizer,
    ) -> peft_trainer.PeftTrainer:
        """Build PeftTrainer (including training config) and attach input callback.

        Args:
            lora_model: LoRA-augmented model to optimize.
            optimizer: Optimizer instance used during training.
            max_steps: Maximum training steps used to shape training config.
            tokenizer: Tokenizer used when building model-input callback.

        Returns:
            Configured ``peft_trainer.PeftTrainer`` instance.
        """
        settings = self.settings
        checkpointing_options = ocp.CheckpointManagerOptions(
            save_interval_steps=settings.save_interval_steps,
            max_to_keep=settings.max_to_keep,
        )
        metrics_logging_options = metrics_logger.MetricsLoggerOptions(
            log_dir=self._tensorboard_root(),
            flush_every_n_steps=20,
        )
        training_config = peft_trainer.TrainingConfig(
            eval_every_n_steps=settings.eval_every_n_steps,
            max_steps=max_steps,
            checkpoint_root_directory=self._checkpoint_root(),
            checkpointing_options=checkpointing_options,
            metrics_logging_options=metrics_logging_options,
        )
        trainer = peft_trainer.PeftTrainer(
            model=lora_model,
            optimizer=optimizer,
            training_config=training_config,
        )
        return trainer.with_gen_model_input_fn(self._gen_model_input_fn(tokenizer))

On the other hand, JAX provides the low-level primitives that enable efficient and hardware-aware distributed training. In particular, this pipeline leverages the device mesh abstraction , which allows precise control over how computation is distributed across TPU cores.

This design enables GemmaEarth to scale from a small number of devices to large TPU configurations without changing the training logic.

Source: https://github.com/haruiz/gemma_earth/blob/5fd2e4c0b6cf3c65eac026bd5aecc136e81072b8/src/gemma_earth/trainers/base.py#L279

def create_mesh(self) -> jax.sharding.Mesh:
  num_devices = jax.local_device_count()
  logger.info("Detected %d local devices for JAX mesh creation.", num_devices)

  if num_devices >= 8:
    return jax.make_mesh(
      (1, 4, 2),
      ("data", "fsdp", "tp"),
      axis_types=(jax.sharding.AxisType.Auto,) * 3,
    )

  if num_devices >= 2:
    return jax.make_mesh(
      (num_devices, 1),
      ("fsdp", "tp"),
      axis_types=(jax.sharding.AxisType.Auto,) * 2,
    )

  return jax.make_mesh(
    (1, 1),
    ("fsdp", "tp"),
    axis_types=(jax.sharding.AxisType.Auto,) * 2,
  )

Model Adaptation (LoRA + Structured State)

Model adaptation is handled through Qwix and Flax NNX , enabling parameter-efficient fine-tuning and structured state management.

Qwix injects LoRA adapters, enabling efficient adaptation without updating the full model
Flax NNX + JAX manages model state and enforces sharding across the device mesh

Result: a LoRA-augmented model with explicitly sharded state , ready for distributed TPU execution.

Source: https://github.com/haruiz/gemma_earth/blob/5fd2e4c0b6cf3c65eac026bd5aecc136e81072b8/src/gemma_earth/trainers/base.py#L328v

def build_lora_model(self, mesh: jax.sharding.Mesh) -> nnx.Module:
  if self.base_model is None:
    raise RuntimeError("Base model was not initialized by load_base_model().")

  settings = self.settings
  lora_provider = qwix.LoraProvider(
    module_path=self.LORA_MODULE_PATH,
    rank=settings.lora_rank,
    alpha=settings.lora_alpha,
  )

  model_input = self.base_model.get_model_input()
  lora_model = qwix.apply_lora_to_model(self.base_model, lora_provider, **model_input)

  with mesh:
    state = nnx.state(lora_model)
    pspecs = nnx.get_partition_spec(state)
    sharded_state = jax.lax.with_sharding_constraint(state, pspecs)
    nnx.update(lora_model, sharded_state)

  return lora_model

Optimization

Optimization is handled via Optax, a JAX-native gradient processing library.

The combination of warmup + cosine decay scheduling with AdamW ensures stable and efficient convergence during LoRA-based fine-tuning, particularly in large-scale TPU environments.

Source: https://github.com/haruiz/gemma_earth/blob/5fd2e4c0b6cf3c65eac026bd5aecc136e81072b8/src/gemma_earth/trainers/base.py#L361

def _build_optimizer(self, max_steps: int) -> optax.GradientTransformation:
  settings = self.settings
  warmup_steps = max(1, int(max_steps * settings.warmup_ratio))
  lr_schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=settings.learning_rate,
    warmup_steps=warmup_steps,
    decay_steps=max(1, max_steps),
    end_value=0.0,
  )

  return optax.adamw(
    learning_rate=lr_schedule,
    weight_decay=settings.weight_decay,
  )

Checkpointing and Experiment Management

Checkpointing is managed using Orbax , which provides flexible and structured model persistence.

Orbax supports:

PyTree-based checkpoints
Partial restoration , allowing recovery even when model structures evolve

This makes experimentation more robust and iterative.

Source: https://github.com/haruiz/gemma_earth/blob/5fd2e4c0b6cf3c65eac026bd5aecc136e81072b8/src/gemma_earth/trainers/base.py#L579

manager = ocp.CheckpointManager(
  self._checkpoint_root(),
  item_handlers={"model_params": ocp.PyTreeCheckpointHandler()},
)
try:
  resolved_step = step if step is not None else manager.latest_step()
  if resolved_step is None:
    raise RuntimeError(f"No checkpoint found under {self._checkpoint_root()}")

  abstract_state = nnx.state(model)
  restore_args = ocp.checkpoint_utils.construct_restore_args(target=abstract_state)

  checkpoint = manager.restore(
    resolved_step,
    args=ocp.args.Composite(
      model_params=ocp.args.PyTreeRestore(
        item=abstract_state,
        restore_args=restore_args,
        partial_restore=True,
      )
    ),
  )

  nnx.update(model, checkpoint.model_params)
  return resolved_step
finally:
  manager.close()

Benchmark results (1.5k samples)

After post-training, for the satellite scene classification task on a 1.5k-sample evaluation dataset, the fine-tuned model shows consistent and significant improvements over the baseline across all metrics, as illustrated in the figure below.

Performance gains are strongest in recall-oriented metrics, while improvements in F1 and Jaccard scores indicate better prediction quality and consistency.

The increase in macro_f1 confirms improvements across both frequent and less frequent classes, and higher exact match scores indicate better full-label-set predictions.

Training dynamics show stable convergence, with evaluation loss decreasing (~0.33 → ~0.22) and perplexity improving (~1.45 → ~1.25), without signs of overfitting. Performance plateaus around 25k–30k steps, suggesting convergence.

More detailed metrics and configurations are available in the GitHub repository.

Future work

The current implementation focuses on the classification subset of EarthDial (BigEarthNet). Future work will extend this pipeline to additional tasks, including captioning, VQA, and reasoning.

Given the modular design, adapting to these tasks primarily involves updating DATASET_RELATIVE_DIR and re-running the pipeline. Detailed instructions are available in the repository.

Closing Thoughts

GemmaEarth demonstrates how the JAX ecosystem can be composed into a practical, production-oriented training stack for domain-specific multimodal models.

Beyond model performance, its key contribution is an extensible engineering blueprint — enabling the transition from experimentation to scalable, reproducible systems leveraging TPU acceleration.

Thanks for taking the time to explore GemmaEarth! If you found this useful, feel free to explore the repository, experiment with the pipeline, or share feedback — I’d love to see how others extend this work.

Happy building!!

Acknowledgements

Google Cloud credits for this project were provided through the #TPUSprint

Although scripts are provided in the GitHub repository to provision GCP services — including TPU setup — you can refer to the official GCP documentation for more detailed guidance on creating and configuring a TPU VM:

https://docs.cloud.google.com/tpu/docs/create-tpu-vm

From Single Agent to Production-Ready Multi Agent Architecture with ADK, MCP, A2A, and AG-UI

Henry Ruiz, PhD — Thu, 19 Feb 2026 21:39:20 +0000

In the previous post, we built an end-to-end single-agent application using ADK and AG-UI. In this post, we evolve that design into a production-ready multi-agent architecture with three coordinated agents:

Weather Agent: retrieves and interprets real-time meteorological data.
Maps Agent: performs geospatial reasoning, including location lookup, coordinate resolution, and map-aware contextual queries.
Atlas Agent: supervises the workflow by delegating tasks to the right specialist, sequencing execution, and synthesizing results into a single, user-ready response.

To power this collaboration, we leverage MCP (Model Context Protocol) to distribute tool-based function execution across agents and A2A (Agent-to-Agent Protocol) as a standardized inter-agent communication layer that ensures consistent, seamless coordination among peers. We further extend the AG-UI (Agent-User Interaction Protocol) integration introduced in the previous tutorial to support these capabilities end-to-end, enabling all three protocols to operate cohesively while preserving each agent’s ability to drive dynamic interface behavior through a unified interaction model. The following figure illustrates the overall architecture of this multi-agent system.

Limitations of Single-Agent Architectures

While single-agent architectures can effectively support many workflows, they encounter fundamental limitations as applications grow in complexity and scope. One of the most prominent challenges is context overload. As tasks become more interdependent and require longer, richer instruction sequences, the amount of information a single agent must process quickly exceeds what is practical or reliable. Recent studies (https://arxiv.org/abs/2307.03172, https://github.com/Arize-ai/LLMTest_NeedleInAHaystack)) show that large language models often exhibit positional bias, placing disproportionate attention on the beginning and end of long prompts while overlooking information in the middle, where critical constraints and task dependencies are commonly specified. This limitation reduces reasoning consistency and degrades performance as task complexity increases.

Image taken from https://arize.com/blog-course/the-needle-in-a-haystack-test-evaluating-the-performance-of-llm-rag-systems/

Distributed Reasoning Through Multi-Agent Systems

A multi-agent system helps address these shortcomings by distributing reasoning across specialized agents, each responsible for a well-defined subset of tasks. Rather than burdening a single agent with managing both the global application state and task handling logic, responsibilities are decomposed across agents with distinct roles, independent context windows, and dedicated toolsets. In doing so, multi-agent systems provide a principled design pattern for enforcing separation of concerns, where reasoning, decision-making, and tool interaction are explicitly isolated across agents rather than entangled within a monolithic control loop. This decomposition reduces cognitive load, improves reasoning reliability, and enables more scalable and maintainable system designs, particularly as task complexity and contextual requirements grow (https://arxiv.org/abs/2305.14325).

When Multi-Agent Architectures Become Necessary

These orchestration needs become especially pronounced in problem settings that exceed the limits of single-agent reasoning and are inherently multidisciplinary. In such scenarios, effective problem solving depends on integrating multiple forms of expertise, heterogeneous reasoning strategies, and specialized tools, each addressing a distinct aspect of the task. This makes coordinated interaction among specialized agents a core system requirement rather than an implementation choice.

Problems that motivate multi-agent architectures typically exhibit four defining characteristics: planning , where solutions require coordinated, multi-step reasoning with explicit dependency management; diverse expertise , where multiple forms of domain-specific reasoning must be applied in concert; extensive context , where relevant information spans large, heterogeneous inputs beyond the reliable capacity of a single agent; and adaptive solutions , where execution paths must evolve dynamically in response to intermediate results, partial failures, or changing constraints. (https://multiagentbook.com/).

Taken together, these properties make multi-agent systems a modular, interpretable, and scalable paradigm for building advanced AI applications. In the following sections, we demonstrate how to transition from a single-agent design to a coordinated multi-agent framework using ADK for agent definition, MCP for tool-based function calling, and A2A for structured agent-to-agent communication.

Building the Multi-Agent System

To evolve our application from a single-agent design into a fully collaborative system, we decompose the main agent’s reasoning into three coordinated specialist agents. The weather_agent is responsible for handling environmental and meteorological queries; the maps_agent focuses on geospatial reasoning and place-based information retrieval; and the atlas_agent integrates and synthesizes information across domains, enabling higher-level reasoning and informed decision-making through cross-agent coordination.

Maps Agent

The Maps Agent is responsible for all geospatial reasoning within the multi-agent architecture. It handles user queries related to location identification, geographic coordinate resolution, and the generation of contextual and semantic descriptions of places. To support these capabilities, we equip the agent with two MCP-integrated tools that combine the Google Maps API with the Gemini API, leveraging Google Maps grounding to (i) resolve precise geographic coordinates and (ii) generate semantically rich, map-aware descriptions of locations. This integration enables the agent to ground generative outputs in authoritative, up-to-date geospatial data while preserving semantic expressiveness.

The MCP server exposes the following tools to the Maps Agent:

get_place_location

Resolves the geographic coordinates (latitude and longitude) for a given place name using the Google Maps Geocoding API.

The tool queries Google Maps directly.
If successful, it returns a structured JSON object containing the latitude and longitude.
If the place cannot be resolved, a descriptive error message is returned.

2. get_place_details

Provides a semantically rich, Gemini-generated description of a place. —

First, the get_place_location tool is invoked to obtain the location coordinates.
It then passes the user’s prompt and the resolved coordinates to the Gemini model, with Google Maps enabled, allowing Gemini to retrieve place-specific context, points of interest, and localized descriptions. —
The result is a high-level narrative or analytic summary generated by Gemini, grounded in the Maps tool’s retrieved information.

Internally, these tools work together to bridge two modalities: Google Maps supplies precise geospatial data. Gemini interprets that data and produces natural-language, context-aware responses.

Here is the implementation of the maps agent and MCP server:

git clone https://github.com/haruiz/atlas_app # clone repo 
cd ./atlas_app # jump into the code folder 
git checkout single_agent # switch to the single_agent branch 
mkdir -p mcp/maps_server # create folder 
touch mcp/maps_server/{server.py,tools.py, __init__.py} # create python files

mcp/maps_server/tools.py

import os
from typing import Any, Dict, Union

import googlemaps
from dotenv import load_dotenv, find_dotenv
from google import genai
from google.genai import types

# Load environment variables
load_dotenv(find_dotenv())

gmaps_client = googlemaps.Client(key=os.getenv("GOOGLE_MAPS_API_KEY"))
genai_client = genai.Client(http_options=types.HttpOptions(api_version="v1"))

# ---------------------------------------------------------------------------
# Utility: Get Geolocation
# ---------------------------------------------------------------------------
def get_place_location(place_name: str) -> Dict[str, Union[str, Dict[str, float]]]:
    """
    Get geographic coordinates from Google Maps API given a place name.
    """
    try:
        geocode_result = gmaps_client.geocode(place_name)
        if not geocode_result:
            return {
                "status": "error",
                "message": f"Could not find coordinates for: {place_name}"
            }
        location = geocode_result[0]["geometry"]["location"]
        return {
            "status": "success",
            "result": {"latitude": location["lat"], "longitude": location["lng"]}
        }

    except Exception as e:
        return {"status": "error", "message": str(e)}

# ---------------------------------------------------------------------------
# Utility: Get Place Details using Gemini + Maps Tool
# ---------------------------------------------------------------------------
def get_place_details(place_name: str, query_prompt: str) -> Dict[str, str]:
    """
    Get place details using Gemini’s Google Maps Tool.
    """
    try:
        # Get coordinates first
        location_result = get_place_location(place_name)

        if location_result["status"] == "error":
            raise Exception(location_result["message"])

        coords = location_result["result"]
        latitude = coords["latitude"]
        longitude = coords["longitude"]

        # Call Gemini with Maps tool
        response = genai_client.models.generate_content(
            model="gemini-2.5-flash",
            contents=query_prompt,
            config=types.GenerateContentConfig(
                tools=[
                    types.Tool(
                        google_maps=types.GoogleMaps(enable_widget=False)
                    )
                ],
                tool_config=types.ToolConfig(
                    retrieval_config=types.RetrievalConfig(
                        lat_lng=types.LatLng(
                            latitude=latitude,
                            longitude=longitude,
                        ),
                        language_code="en_US",
                    )
                ),
            ),
        )

        return {
            "status": "success",
            "result": response.text
        }

    except Exception as e:
        return {
            "status": "error",
            "message": f"Failed to get place details: {str(e)}"
        }

mcp/maps_server/server.py

from typing import Any

from mcp.server.fastmcp import FastMCP

from tools import get_place_location, get_place_details

# Create an MCP server
mcp = FastMCP("geoloc_server", port=8080)

@mcp.tool(name="get_place_location", description="Get coordinates from an address using Google Maps API.")
async def get_place_location_tool(place_name: str) -> dict[str, Any]:
    """
    Get coordinates from an address using Google Maps API.
    :param place_name:
    :return:
    """
    return get_place_location(place_name)

@mcp.tool(name="get_place_details", description="Get place details using Google Maps Tool in Gemini.")
async def get_place_details_tool(place_name: str, query_prompt: str) -> dict[str, Any]:
    """
    Get place details using Google Maps Tool in Gemini.
    :param query_prompt: User query prompt for place details
    :param place_name: Name of the place to get details for
    :return:
    """
    return get_place_details(place_name, query_prompt)

if __name__ == " __main__":
    mcp.run(transport="streamable-http")

Weather Agent

The Weather Agent is responsible for handling all user requests related to real-time atmospheric conditions and environmental insights. It provides structured meteorological information — such as temperature, humidity, wind speed, apparent temperature, and qualitative weather conditions — based on geospatial coordinates. To support these capabilities, the agent leverages MCP-accessible tools backed by the Open‑Meteo API for authoritative weather data retrieval and the Gemini API for reasoning, interpretation, and contextual synthesis. This combination enables the Weather Agent to deliver accurate, location-aware environmental insights while maintaining coherent, user-friendly responses.

The MCP server exposes the following tool to the Weather Agent:

get_weather

Retrieves current weather conditions for a specific latitude and longitude by querying the Open-Meteo API. The tool processes raw atmospheric measurements (temperature, humidity, wind speed, gust levels, and weather code) and maps weather codes to human-readable conditions using a predefined lookup table. The output is returned as a structured JSON object containing both the raw values and their semantic interpretation.

Internally, the tool workflow includes:

Fetching meteorological variables via an asynchronous HTTP request.
Translating numerical weather codes into descriptive categories (e.g., Clear sky, Moderate rain, Fog).
Packaging the processed data into a standardized result format compatible with MCP.

Below is the code for the weather agent and MCP server.

mkdir -p mcp/weather_server # create folder 
touch mcp/weather_server/{server.py,tools.py, __init__.py} # create python files

mcp/weather_server/tools.py

import httpx
from typing import Dict, Any, Optional

# ---------------------------------------------------------------------------
# Weather Code Mapping (constant)
# ---------------------------------------------------------------------------
WEATHER_CODE_MAP: Dict[int, str] = {
    0: "Clear sky",
    1: "Mainly clear",
    2: "Partly cloudy",
    3: "Overcast",
    45: "Fog",
    48: "Depositing rime fog",
    51: "Light drizzle",
    53: "Moderate drizzle",
    55: "Dense drizzle",
    61: "Light rain",
    63: "Moderate rain",
    65: "Heavy rain",
    71: "Light snow",
    73: "Moderate snow",
    75: "Heavy snow",
    80: "Rain showers",
    95: "Thunderstorm",
    99: "Thunderstorm with hail",
}

def get_weather_condition(code: Optional[int]) -> str:
    """
    Convert a weather code into a human-readable description.

    Args:
        code: Weather condition code.

    Returns:
        A descriptive string representing the weather condition.
    """
    return WEATHER_CODE_MAP.get(code, "Unknown")

# ---------------------------------------------------------------------------
# Weather Fetching
# ---------------------------------------------------------------------------
async def get_weather(latitude: float, longitude: float) -> Dict[str, Any]:
    """
    Fetch current weather conditions for a given latitude and longitude.

    Args:
        latitude: Latitude coordinate.
        longitude: Longitude coordinate.

    Returns:
        A dictionary containing:
            - status: "success" or "error"
            - result: weather fields (if success)
            - message: error message (if error)
    """
    url = (
        "https://api.open-meteo.com/v1/forecast?"
        f"latitude={latitude}&longitude={longitude}"
        "&current=temperature_2m,apparent_temperature,relative_humidity_2m,"
        "wind_speed_10m,wind_gusts_10m,weather_code"
    )

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            response = await client.get(url)
            response.raise_for_status()

            data = response.json()
            current = data.get("current")

            if not current:
                return {"status": "error", "message": "Weather API returned no 'current' data."}

            result = {
                "temperature": current.get("temperature_2m"),
                "feelsLike": current.get("apparent_temperature"),
                "humidity": current.get("relative_humidity_2m"),
                "windSpeed": current.get("wind_speed_10m"),
                "windGust": current.get("wind_gusts_10m"),
                "conditions": get_weather_condition(current.get("weather_code")),
            }

            return {"status": "success", "result": result}

    except httpx.HTTPError as e:
        return {"status": "error", "message": f"HTTP error: {str(e)}"}

    except Exception as e:
        return {"status": "error", "message": str(e)}

# ---------------------------------------------------------------------------
# Example Execution
# ---------------------------------------------------------------------------
if __name__ == " __main__":
    import asyncio

    # Example: New York City coordinates
    nyc_lat, nyc_lon = 40.7128, -74.0060

    weather = asyncio.run(get_weather(nyc_lat, nyc_lon))
    print("Weather:", weather)

mcp/weather_server/server.py

from typing import Any

from mcp.server.fastmcp import FastMCP

from tools import get_weather

# Create an MCP server
mcp = FastMCP("weather_server", port=8081, json_response=True)

@mcp.tool(name="get_weather", description="Get weather information for a given location.")
async def get_weather_tool(latitude: float, longitude: float) -> dict:
    """
    Get the weather for a given location.
    :param latitude: Latitude of the location.
    :param longitude: Longitude of the location.
    :return:
    """
    return await get_weather(latitude, longitude)

if __name__ == " __main__":
    mcp.run(transport="streamable-http")

The instructions defining the behavior of both the maps and weather agents are available in the GitHub repository accompanying this blog post, in the multiple_agent branch.

GitHub - haruiz/atlas_app at multiple_agent

Atlas Agent

In the single-agent example, the Atlas Agent was responsible for managing all reasoning and decision-making within a single control loop. In this tutorial, however, the Atlas Agent transitions into a pure orchestration role, delegating domain-specific reasoning to specialized agents and focusing on coordination rather than execution. This section highlights the key changes required to evolve the Atlas Agent from a monolithic reasoning component into a multi-agent orchestrator.

Implemented using ADK(Agent Development Kit), the Atlas Agent supports A2A-based task delegation, optional debugging callbacks, and memory preloading. Its instruction set defines the end-to-end execution flow, including location resolution, weather data requests, sequencing of agent and tool calls, and the integration of multi-agent outputs into clear, user-friendly summaries.

Next, we implement this orchestration logic in code.

pycharm ./agents/atlas_agent/agent.py

agents/atlas_agent/agent.py

from typing import Optional, Dict, Any

from a2a.utils import AGENT_CARD_WELL_KNOWN_PATH
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.llm_agent import Agent
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
from google.adk.models import LlmRequest, LlmResponse
from google.adk.tools import BaseTool, ToolContext, AgentTool
from google.adk.tools.preload_memory_tool import PreloadMemoryTool
from dotenv import load_dotenv, find_dotenv

# Load environment variables
load_dotenv(find_dotenv())

# ======================================================================
# CALLBACKS
# ======================================================================

def before_model(
    callback_context: CallbackContext,
    llm_request: LlmRequest,
) -> Optional[LlmResponse]:
    """
    Executes before the model handles a request.
    Allows inspection of inputs, debugging, or intercepting execution.

    Return:
        - LlmResponse to skip model execution
        - None to allow normal execution
    """
    agent_name = callback_context.agent_name
    print(f"[Orchestrator Callback] before_model for agent: {agent_name}")

    last_user_message = ""
    if llm_request.contents and llm_request.contents[-1].role == "user":
        parts = llm_request.contents[-1].parts
        if parts:
            last_user_message = parts[0].text

    print(f"[Orchestrator Callback] User message: {last_user_message}")

    print("[Orchestrator Callback] Tools available to model:")
    for tool in llm_request.config.tools or []:
        print(f" - {tool.name}")

    return None

def before_tool(
    tool: BaseTool, args: Dict[str, Any], tool_context: ToolContext
) -> Optional[Dict]:
    """
    Executes before any tool call.
    Allows modification of arguments or skipping execution.
    """
    print(
        f"[Orchestrator Callback] before_tool: '{tool.name}' "
        f"(agent: {tool_context.agent_name})"
    )
    print(f"[Orchestrator Callback] Tool arguments: {args}")
    return None

# ======================================================================
# REMOTE AGENTS (enable when servers are available to test a2a locally)
# ======================================================================
# maps_agent = RemoteA2aAgent(
# name="maps_agent",
# description="Provides geocoding, place lookup, and mapping utilities.",
# agent_card=f"http://127.0.0.1:8001{AGENT_CARD_WELL_KNOWN_PATH}",
# )
#
# weather_agent = RemoteA2aAgent(
# name="weather_agent",
# description="Retrieves real-time weather conditions for a given coordinate.",
# agent_card=f"http://127.0.0.1:8002{AGENT_CARD_WELL_KNOWN_PATH}",
# )

# ======================================================================
# ORCHESTRATOR AGENT
# ======================================================================

root_agent = Agent(
    name="atlas_agent",
    model="gemini-2.5-pro",

    # Uncomment for debugging
    # before_tool_callback=before_tool,
    # before_model_callback=before_model,

    tools=[
        PreloadMemoryTool(),

        # When remote agents are active:
        # AgentTool(agent=maps_agent),
        # AgentTool(agent=weather_agent),
    ],

    instruction="""
    You are the Orchestrator Agent. Your role is to coordinate two specialized agents:

    - maps_agent: Resolves place names, coordinates, and map-related queries.
    - weather_agent: Retrieves weather information using geographic coordinates.

    Your responsibility is to determine when each agent should be called, how their outputs should be used, and how to produce a final integrated response for the user.

    ----------------------------------------------------------------------
    Workflow for Weather Queries
    ----------------------------------------------------------------------

    When the user requests weather information for a location:

    1. Call maps_agent  
       Input: the place name or location text.  
       Output: latitude and longitude.

    2. Call weather_agent  
       Input: coordinates returned by maps_agent.  
       Output: structured weather information.

    3. Combine both results into a clear, user-facing response.

    ----------------------------------------------------------------------
    Workflow for Location or Mapping Queries
    ----------------------------------------------------------------------

    If the user asks only for location information, directions, or coordinates:

    1. Call maps_agent directly.  
    2. Return the result to the user.

    ----------------------------------------------------------------------
    Execution Rules
    ----------------------------------------------------------------------

    - Execute agent calls in sequential order.  
    - Do not call multiple agents at the same time.  
    - Always wait for the response of the previous agent before making a new call.  
    - Pass intermediate outputs to the next agent when required.  
    - Do not fabricate or assume agent outputs; always rely on actual tool responses.  
    - Provide a final, coherent response after completing the workflow.

    ----------------------------------------------------------------------
    """,
)

MCP, A2A, and AG-UI Integration

Previously, we demonstrated how MCP standardizes interactions between models, agents, and tools, enabling a clear separation between the tools layer and agent implementations. This architectural decoupling improves scalability and maintainability while ensuring the tools layer remains fully language-agnostic, allowing tools to be implemented in any programming language and executed independently of the agent’s runtime environment.

Building on this foundation, we extend the architecture by transforming our ADK agents into remote agents using the Agent-to-Agent (A2A) protocol. A2A enables agents to communicate across network boundaries and operate as independently deployable services, significantly enhancing modularity, fault isolation, and scalability in distributed environments.

Beyond communication, A2A also brings structure. It standardizes how agent capabilities are defined and exposed, how requests are routed, and how messages and responses flow between agents. As a result, agents can collaborate as loosely coupled services that are easy to reason about, evolve, and operate.

The payoff is significant: agents can now be developed, scaled, updated, and even recovered independently, dramatically improving fault tolerance, system resilience, and long-term maintainability in real-world multi-agent systems.

Exposing ADK Agents via A2A

To begin, we need to convert our existing ADK agents into A2A-compliant endpoints so they can communicate with one another. This is the core step behind the question:

I have an agent — how do I expose it so that other agents can use it via A2A?

This capability is essential for building multi-agent systems where agents collaborate, delegate tasks, and exchange information programmatically. ADK offers two primary methods for exposing an agent through A2A:

Using to_a2a(root_agent)

This is the simplest path when you want to convert an existing agent into an A2A endpoint. It gives you full control over deployment — e.g., exposing the agent via uvicorn rather than relying on adk deploy api_server. The to_a2a() function also auto-generates the agent card based on your code.

Creating and hosting your own agent.json card via adk api_server — a2a

This approach integrates seamlessly with adk web for debugging and testing. It also allows you to host multiple agents from a single parent directory, each exposed automatically as long as an agent.json file is present. The tradeoff is that you must author the agent card manually (see the A2A Python tutorial for details).

Agent Development Kit

Approach Used in This Tutorial

For this tutorial, we will use the to_a2a() function, as it provides the most direct and automated path for exposing an ADK agent as an A2A endpoint while generating the agent card behind the scenes. If you’re interested in the alternative adk api_server approach, refer to the full ADK documentation for additional guidance.

Let’s create and copy the following file to the weather_agent and maps_agent folders:

agents/maps_agent/a2a_server.py

from google.adk.a2a.utils.agent_to_a2a import to_a2a
from agent import root_agent

# Make your agent A2A-compatible
a2a_app = to_a2a(root_agent, port=8001)

if __name__ == ' __main__':
    import uvicorn
    uvicorn.run("a2a_server:a2a_app", host="127.0.0.1", port=8001, reload=True, workers=1)

agents/weather_agent/a2a_server.py

from google.adk.a2a.utils.agent_to_a2a import to_a2a
from agent import root_agent

# Make your agent A2A-compatible
a2a_app = to_a2a(root_agent, port=8002)

if __name__ == ' __main__':
    import uvicorn
    uvicorn.run("a2a_server:a2a_app", host="127.0.0.1",port=8002, reload=True, workers=1)

Now that we have our MCP servers in place and our A2A-compliant ADK agents running, the next step is to update the main application entry point. Specifically, we will modify the backend so the Atlas Agent is wrapped as an AG-UI–compatible agent, building directly on the previous example.

At this point, it is important to highlight an intentional design choice: during the implementation of the atlas_agent, we deliberately did not enable the remote agents, nor did we register them as sub-agents within the orchestrator’s definition. This is because, in this architecture, A2A task delegation and inter-agent messaging are managed at the frontend layer, within the Next.js application. The Next.js frontend acts as the coordination surface for user requests originating from the AG-UI interface, controlling how those requests are dispatched, tracked, and observed across the multi-agent system.

To support this flow, we need to update two key components:

The backend main.py file, where the agent is exposed in an AG-UI–compatible form
The frontend route.ts file in our Next.js application, which handles incoming user requests and streams agent updates to the UI

The core reason for these changes is that AG-UI must be able to intercept messages produced by the A2A-based multi-agent architecture. By doing so, AG-UI can observe agent interactions in real time and stream structured updates back to the React frontend, enabling a responsive and transparent user experience.

In the next sections, we’ll walk through the required updates on both the backend and frontend to make this integration work end-to-end.

First, let’s update the backend entry point: main.py. This is where we’ll wrap the Atlas Agent as an AG-UI–compatible agent and expose the endpoint that the frontend will connect to.

main.py

"""
Orchestrator Agent - Coordinates between Weather and Maps agents.
Communicates using the AG-UI Protocol for UI integration, and delegates tasks
to specialized A2A agents through ADK middleware.
"""

from __future__ import annotations

# Load environment variables (API keys, configs, etc.)
from dotenv import load_dotenv

load_dotenv()

import uvicorn
from fastapi import FastAPI

# ADK wrapper that exposes agents through the AG-UI Protocol
from ag_ui_adk import ADKAgent, add_adk_fastapi_endpoint
from agents.atlas_agent.agent import root_agent as orchestrator_agent

# -------------------------------------------------------------------
# Wrap the LlmAgent using ADKAgent
# This adds:
# - session handling
# - stateful memory
# - tool routing
# - AG-UI Protocol support
# -------------------------------------------------------------------

adk_orchestrator_agent = ADKAgent(
    adk_agent=orchestrator_agent, # The LLM agent defined above
    app_name="orchestrator_app", # ID for the AG-UI front-end
    user_id="demo_user", # Session tracking user ID
    session_timeout_seconds=3600, # Auto-expire session after 1 hour
    use_in_memory_services=True # Enables in-memory memory + tool services
)

# -------------------------------------------------------------------
# FastAPI application setup
# Adds the ADK Orchestrator as an API endpoint using AG-UI Protocol
# -------------------------------------------------------------------

app = FastAPI(title="A2A Orchestrator (ADK + AG-UI Protocol)")

# Inject the orchestrator endpoint at "/"
add_adk_fastapi_endpoint(app, adk_orchestrator_agent, path="/")

# -------------------------------------------------------------------
# Run the development server (only when executing `python main.py`)
# Uvicorn reloads on changes and ensures ADK's MCP tools run properly
# -------------------------------------------------------------------

if __name__ == ' __main__':
    uvicorn.run(
        "main:app",
        host="localhost",
        port=8000,
        reload=True, # Auto-reload server on code changes (dev mode)
        workers=1 # ADK/MCP tools require single-worker execution
    )

Next, we’ll move to the frontend and update ui/app/api/copilotkit/route.ts. This route will act as the bridge between the AG-UI interface and our remote multi-agent backend, handling requests and streaming agent updates back to the UI.

ui/app/api/copilotkit/route.ts

import { A2AMiddlewareAgent } from "@ag-ui/a2a-middleware"
import {CopilotRuntime, copilotRuntimeNextJSAppRouterEndpoint, ExperimentalEmptyAdapter} from "@copilotkit/runtime";
import {AbstractAgent, HttpAgent} from "@ag-ui/client";
import {NextRequest} from "next/server";

export async function POST(request: NextRequest) {
// These first two are the urls to the a2a agents
const weatherAgentUrl = process.env.WEATHER_AGENT_URL || "http://localhost:8002";
const mapsAgentUrl = process.env.MAPS_AGENT_URL || "http://localhost:8001";
const orchestratorUrl = process.env.ORCHESTRATOR_URL || "http://localhost:8000";

// the orchestrator agent we pass to the middleware needs to be an instance of a derivative of an ag-ui `AbstractAgent`
// In this case, we have access to the agent via url, so we can gain an instance using the `HttpAgent` class
const orchestrationAgent: AbstractAgent = new HttpAgent({
  url: orchestratorUrl,
});

// A2A Middleware: Wraps orchestrator and injects send_message_to_a2a_agent tool
// This allows orchestrator to communicate with A2A agents transparently
const a2aMiddlewareAgent = new A2AMiddlewareAgent({
  description:
    "An orchestrator agent that coordinates two specialized agents: Weather Agent and Maps Agent build on the ADK framework.",
  // We pass the urls to the a2a agents, the middleware will handle the connections
  agentUrls: [
    weatherAgentUrl,
    mapsAgentUrl,
  ],
  // Pass the agent instance
  orchestrationAgent
});

// CopilotKit runtime connects frontend to agent system
  // CopilotKit runtime connects frontend to agent system
  const runtime = new CopilotRuntime({
    agents: {
      a2a_chat: a2aMiddlewareAgent, // Must match agent prop in <CopilotKit agent="a2a_chat">
    },
  });

  const { handleRequest } = copilotRuntimeNextJSAppRouterEndpoint({
    runtime,
    serviceAdapter: new ExperimentalEmptyAdapter(),
    endpoint: "/api/copilotkit",
  });

  return handleRequest(request);
}

The most important piece of this frontend/Nextjs App integration is the A2A middleware layer, implemented via A2AMiddlewareAgent. Conceptually, it sits between your AG-UI/CopilotKit runtime and your remote A2A agents, turning a distributed set of services into a single, seamless “agent” the UI can talk to.

In this implementation, the middleware wraps the orchestrator, which is exposed as an AG UI AbstractAgent via HttpAgent, and registers the URLs of the specialized remote agents for Maps and Weather. From there, it injects a standardized tool interface, most notably a send_message_to_a2a_agent capability, allowing the orchestrator(atlas_agent) to delegate work to remote agents without having to manage low-level networking, request routing, or message formatting directly.

In other words, the orchestrator continues to focus on decision-making and sequencing, while the middleware handles the heavy lifting required for remote collaboration:

Agent connectivity: manages connections to multiple A2A agent endpoints via their URLs
Request routing: forwards orchestrator calls to the appropriate specialized agent
Message exchange + coordination: standardizes how messages are sent and responses are returned

This is why, in the code, the UI only needs to register a single runtime agent:

agents: {
  a2a_chat: a2aMiddlewareAgent,
}

We finally need to update our page.tsx file so the UI can properly render A2A interactions. The complete implementations of the MessageToA2A and MessageFromA2A components are available in the GitHub repository under: atlas_app, multiple_agent branch:

atlas_app/ui/components/a2a at multiple_agent · haruiz/atlas_app

ui/app/page.tsx

"use client";

import React, {useEffect} from "react";
import "@copilotkit/react-ui/styles.css";
// import "./style.css";

// CopilotKit core
import {useCopilotAction, useCopilotChat,} from "@copilotkit/react-core";
import {CopilotChat} from "@copilotkit/react-ui";
import {MessageToA2A} from "@/components/a2a/MessageToA2A";
import {MessageFromA2A} from "@/components/a2a/MessageFromA2A";

const Chat = () => {
    const {visibleMessages} = useCopilotChat();

    useCopilotAction({
    name: "send_message_to_a2a_agent",
    description: "Sends a message to an A2A agent",
    available: "frontend",
    parameters: [
      {
        name: "agentName",
        type: "string",
        description: "The name of the A2A agent to send the message to",
      },
      {
        name: "task",
        type: "string",
        description: "The message to send to the A2A agent",
      },
    ],
    render: (actionRenderProps) => {
        console.log(actionRenderProps);
      return (
        <>
          <MessageToA2A {...actionRenderProps} />
          <MessageFromA2A {...actionRenderProps} />
        </>
      );
    },
  });

    useEffect(() => {
      const extractDataFromMessages = () => {
      for (const message of visibleMessages) {
        const msg = message;

        if (msg.type === "ResultMessage" && msg.actionName === "send_message_to_a2a_agent") {

          try {
            let parsed;
            const result = msg.result;
            if (typeof result === "string") {
              try {
                  parsed = JSON.parse(result);
              }
            catch {
                continue;
            }
            } else if (typeof result === "object") {
                parsed = result;
            }
            console.log("Sending message to A2A agent: ", parsed);
          } catch (e) {
            console.error("Failed to extract data from message:", e);
          }
        }
      }
    };

    extractDataFromMessages();
    }, [visibleMessages]);

    /* --------------------------------------------------------------------------------------------
     * MAIN UI RENDERING
     * ------------------------------------------------------------------------------------------*/
    return (
        <div
            className="flex justify-center items-center h-full w-full"
            data-testid="background-container"
        >
            <div className="h-full w-full md:w-8/10 md:h-8/10 rounded-lg">
                <CopilotChat
                    onSubmitMessage={(message) => {
                        console.log("User submitted message:", message);
                    }}
                    suggestions={[{
                        title: "What is the weather in San Francisco?",
                        message: "What is the weather in San Francisco?"
                    }]}
                    className="h-full rounded-2xl max-w-6xl mx-auto"
                    labels={{initial: "Hi, I'm an agent. Want to chat?"}}
                />
            </div>
            <span className="absolute bottom-4 right-4 text-xs text-gray-400">
                Num Messages: {visibleMessages.length}
            </span>
        </div>
    );
};

/* ------------------------------------------------------------------------------------------------
 * HOME PAGE WRAPPER
 * ----------------------------------------------------------------------------------------------*/

const HomePage = () => {
    return (
        <main className="h-screen w-screen">
            <Chat/>
        </main>
    );
};

export default HomePage;

Before wrapping up, here is a short video demo of the system in action, showing how user requests flow through the AG-UI frontend, are routed via the A2A middleware, and coordinated across the Maps, Weather, and Atlas agents in real time.

You now have a complete end-to-end multi-agent system that is modular, remotely deployable, and easy to extend. MCP keeps your tools cleanly separated and language-agnostic, A2A turns each ADK agent into a network-accessible service that can collaborate with peers, and AG-UI makes those interactions visible and interactive in the frontend. The result is an architecture where you can add new specialist agents, swap tool implementations, or scale components independently without rewriting your application.

From here, try adding a third agent, such as traffic, air quality, or travel planning, and update the A2A middleware instructions to route requests dynamically. As your agent ecosystem grows, you will start to see the real payoff of this approach: clearer separation of concerns, better reliability under complexity, and a user experience that remains transparent even as the system becomes increasingly distributed.

I hope this tutorial proves useful as you explore and build your own multi-agent systems. Whether you are experimenting with new agent interactions or designing production-ready architectures, I hope the patterns and examples shared here help accelerate your work. Thanks for following along, and I look forward to seeing you in the next post.

Note: The full, fully functional implementation of the application is available in the GitHub repository under the multiple_agent branch. To run the complete app end-to-end, use the provided run_app.sh script. This script automatically dispatches all required components—including the MCP servers, the A2A agents, the backend API, and the frontend—so the entire system can be launched with a single command.

Acknowledgements

This work was supported by Google ML Developer Programs and the Google Developers Program , which provided Google Cloud credits and high-quality technical resources #AISprint

Building interactive agentic applications using ADK and AG-UI protocol

Henry Ruiz, PhD — Sat, 20 Dec 2025 18:31:18 +0000

The rise of multi-agent systems is reshaping how we design and deploy modern AI software. As LLMs become better at reasoning and planning, we are moving beyond “one prompt, one response” and toward applications where agents can take actions, call tools, and collaborate to automate complex, real-world workflows (see image below). However, as these systems grow in capability and complexity, a practical bottleneck quickly emerges: it is not only about making the model smarter, but about making the system usable and user-friendly.

From standalone LLMs to agentic systems. A comparison between a single-prompt LLM interaction and a multi-agent system where agents combine reasoning, memory, tools, and retrieval to take actions in an environment, receive feedback, and collaborate with humans, APIs, and external tools.

In practice, many of these usability challenges arise at the interaction layer between backend agent execution and the frontend user interface. Multi-agent systems are inherently multi-turn, asynchronous, and stateful, and without a well-defined interface contract, both users and developers quickly lose visibility and control over what the system is doing. As a result, several recurring challenges consistently appear:

Complex conversation state : user intent, agent context, and partially completed tasks must be preserved and evolved across turns, rather than implicitly reconstructed from prompt history.
Coordination overhead : multiple agents operate in parallel and respond asynchronously; naïve chat-based UIs collapse these interactions into an unreadable stream.
Latency and user experience : multi-step reasoning and external tool calls introduce delays that must be communicated as progress, not experienced as silence.
Conflict resolution : specialized agents may disagree, requiring clear mechanisms for comparison, arbitration, and guidance.
Transparency and attribution : users need to know which agent produced which output in order to maintain trust.
Human-in-the-loop control : real-world systems require intervention points to approve, reject, pause, or refine agent actions during execution.
Scalability of interaction : as the number of agents and the depth of interaction grow, cognitive load becomes the dominant failure mode.

AG-UI is designed to address many of these challenges by standardizing the connection between agent execution and the user interface. It introduces an event-driven interaction protocol that allows agents to emit structured, typed events — such as messages, tool calls, tool results, state updates, and lifecycle signals — that the frontend can interpret deterministically and render consistently.

Image source: https://www.copilotkit.ai/blog/introducing-ag-ui-the-protocol-where-agents-meet-users

By externalizing agent execution as observable, structured events, AG-UI makes conversation state explicit across turns. By attaching agent identity and role metadata to each event, it enables transparency and supports conflict resolution. By streaming execution status and intermediate signals, it mitigates latency through continuous, real-time feedback. Finally, by supporting UI-addressable actions, AG-UI enables human-in-the-loop control without coupling UI logic to agent prompts.

In this post, we will walk through how to integrate AG-UI with ADK to build engaging, agentic user experiences. As a concrete example, we will develop a system that provides users with both weather information and geographic context. To support real-time interactions, our agents will be equipped with tools that retrieve live weather data and interact with the Google Maps API.

Let’s get started.

Create ADK Agent

First, we’ll create our uv project and install the necessary dependencies. Let’s call our application "atlas". To do so, run the following commands in the PyCharm terminal:

uv init atlas --app # Create uv project
cd atlas # move into project folder
uv add ag-ui-adk google-adk google-genai googlemaps # install dependecies

pycharm ./ # Open project on code editor (pycharm)

Next, we create our atlas_agent:

mkdir agents # create agents folder
cd agents # move into agents folder
adk create atlas_agent # create an agent using adk
pycharm atlas_agent/agent.py # open agent file in the editor

Create/update the content of the following files:

atlas_agent/agent.py

from google.adk.agents.llm_agent import Agent
from google.adk.tools.preload_memory_tool import PreloadMemoryTool
from .tools import get_weather, get_place_location, get_place_details

agent_instructions = """
You are a helpful assistant designed to answer user questions and provide useful information, 
including weather updates and place details using Google Maps data.

Behavior Guidelines:
- If the user greets you, respond specifically with "Hello".
- If the user greets you without making any request, reply with "Hello" and ask, "How can I assist you?"
- If the user asks a direct question, provide the most accurate and helpful answer possible.

Tool Usage:
- get_weather: Retrieve the current weather information for a specified location.
- get_place_location: Obtain the precise latitude and longitude of a specified place.
- get_place_details: Fetch detailed information about a place using its geographic coordinates.

Always choose the most appropriate tool to fulfill the user's request, and respond clearly and concisely.
"""
root_agent = Agent(
    name="assistant", # Internal agent name
    model="gemini-2.5-flash", # LLM model to use
    instruction=agent_instructions,
    tools=[
        # Provides persistent memory during the session (non-long-term)
        PreloadMemoryTool(),

        # Direct tool integration example
        get_weather,
        get_place_location,
        get_place_details,

    ]
)

atlas_agent/tools.py

import os
from typing import Optional, Any

import googlemaps
import httpx
from google import genai
from google.genai import types
from dotenv import load_dotenv, find_dotenv

load_dotenv(find_dotenv())

gmaps_client = googlemaps.Client(key=os.getenv("GOOGLE_MAPS_API_KEY"))
genai_client = genai.Client(http_options=types.HttpOptions(api_version="v1"))

def get_weather_condition(code: int) -> str:
    """Map weather code to human-readable condition.

    Args:
        code: WMO weather code.

    Returns:
        Human-readable weather condition string.
    """
    conditions = {
        0: "Clear sky",
        1: "Mainly clear",
        2: "Partly cloudy",
        3: "Overcast",
        45: "Foggy",
        48: "Depositing rime fog",
        51: "Light drizzle",
        53: "Moderate drizzle",
        55: "Dense drizzle",
        56: "Light freezing drizzle",
        57: "Dense freezing drizzle",
        61: "Slight rain",
        63: "Moderate rain",
        65: "Heavy rain",
        66: "Light freezing rain",
        67: "Heavy freezing rain",
        71: "Slight snow fall",
        73: "Moderate snow fall",
        75: "Heavy snow fall",
        77: "Snow grains",
        80: "Slight rain showers",
        81: "Moderate rain showers",
        82: "Violent rain showers",
        85: "Slight snow showers",
        86: "Heavy snow showers",
        95: "Thunderstorm",
        96: "Thunderstorm with slight hail",
        99: "Thunderstorm with heavy hail",
    }
    return conditions.get(code, "Unknown")

async def get_weather(location: str) -> dict:
    """Get the weather for a given location.

    Args:
        location: City name.

    Returns:
        Dictionary with weather information including temperature, feels like,
        humidity, wind speed, wind gust, conditions, and location name.
    """
    try:
        async with httpx.AsyncClient() as client:
            # Geocode the location
            geocoding_url = (
                f"https://geocoding-api.open-meteo.com/v1/search?name={location}&count=1"
            )
            geocoding_response = await client.get(geocoding_url)
            geocoding_data = geocoding_response.json()
            print(geocoding_data)

            if not geocoding_data.get("results"):
                raise ValueError(f"Location '{location}' not found")

            result = geocoding_data["results"][0]
            latitude = result["latitude"]
            longitude = result["longitude"]
            name = result["name"]

            # Get weather data
            weather_url = (
                f"https://api.open-meteo.com/v1/forecast?"
                f"latitude={latitude}&longitude={longitude}"
                f"&current=temperature_2m,apparent_temperature,relative_humidity_2m,"
                f"wind_speed_10m,wind_gusts_10m,weather_code"
            )
            weather_response = await client.get(weather_url)
            weather_data = weather_response.json()

            current = weather_data["current"]

            result = {
                "temperature": current["temperature_2m"],
                "feelsLike": current["apparent_temperature"],
                "humidity": current["relative_humidity_2m"],
                "windSpeed": current["wind_speed_10m"],
                "windGust": current["wind_gusts_10m"],
                "conditions": get_weather_condition(current["weather_code"]),
                "location": name,
            }
            return {
                "status": "success",
                "result": result
            }
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

def get_place_location(place_name: str) -> dict[str, str]:
    """Get coordinates from an address using Google Maps API.

    Args:
        place_name: The name of the place to get coordinates for.

    Returns:
        A dictionary with the status of the operation and the result.
        If successful, the result contains the latitude and longitude.
    """
    try:
        geocode_result = gmaps_client.geocode(place_name)
        if geocode_result is None:
            return {
                "status": "error",
                "message": f"Could not find coordinates for address: {place_name}"
            }

        location = geocode_result[0]["geometry"]["location"]
        lat = location["lat"]
        lng = location["lng"]
        return {
            "status": "success",
            "result": {"latitude": lat, "longitude": lng}
        }
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

def get_place_details(query_prompt: str, latitude: float, longitude: float, model_name: Optional[str]= "gemini-2.5-flash") -> dict[str, Any]:
    """Get place details using Google Maps Tool in Gemini.

    Args:
        query_prompt: The prompt to search for.
        latitude: The latitude of the location.
        longitude: The longitude of the location.
        model_name: The name of the model to use.

    Returns:
        A dictionary with the status of the operation and the result.
        If successful, the result contains the place details.
    """
    try:
        response = genai_client.models.generate_content(
            model=model_name,
            contents=query_prompt,
            config=types.GenerateContentConfig(
                tools=[
                    # Use Google Maps Tool
                    types.Tool(google_maps=types.GoogleMaps(
                        enable_widget=False # Optional: return Maps widget token
                    ))
                ],
                tool_config=types.ToolConfig(
                    retrieval_config=types.RetrievalConfig(
                        lat_lng=types.LatLng( # Pass geo coordinates for location-aware grounding
                            latitude=latitude,
                            longitude=longitude,
                        ),
                        language_code="en_US", # Optional: localize Maps results
                    ),
                ),
            ),
        )
        return {
            "status": "success",
            "result": response.text
        }
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

To see the agent in action, you can use the ADK web UI. It offers a clean, built-in environment for testing, debugging, and experimenting with your agent. Just run the command below, and a browser window will open showing a UI similar to Image 2.

adk web --reload

ADK Web UI

2. Building the AG-UI frontend

Even though we can test our agent within the built-in ADK UI, real-world agentic applications often require a dedicated frontend capable of managing streaming outputs, coordinating asynchronous agent tasks, rendering custom UI components, preserving conversation state, and supporting human-in-the-loop control. AG-UI addresses these architectural needs by providing a standardized protocol for building frontends for multi-agent applications. It defines how agent state, UI intents, and user interactions flow between your model or agent runtime and your application’s frontend — allowing you to ship reliable, debuggable, and user-friendly agentic features quickly. In this section, we will learn how to build the frontend of our atlas application by leveraging AG-UI’s out-of-the-box integration with ADK to create a seamless, interactive agent experience.

2.1. Create Next.js Application

We start creating a Next.js App with the default settings, running the command:

npx create-next-app@latest ui

The command should have created a ui folder with the following structure:

cd ./ui
tree -L 2

├── README.md
├── app
├── eslint.config.mjs
├── next-env.d.ts
├── next.config.ts
├── node_modules
├── package-lock.json
├── package.json
├── postcss.config.mjs
├── public
└── tsconfig.json

Then, we check if our Next.js project was successfully created by running our application in the browser with the command:

yarn dev

Next js App running in the browser

If everything worked correctly in the previous step, and our application is open in the browser, we can now install the AG-UI dependencies and start building our frontend.

yarn add @copilotkit/react-ui @copilotkit/react-core @copilotkit/runtime @ag-ui/client

2.2. Standardizing Backend-to-Frontend Communication with AG-UI

As described earlier, AG-UI provides a standardized communication layer between the agent runtimes (such as ADK, LangGraph, or custom orchestrators) and the frontend. It defines:

How the agent state is exchanged with the frontend
how UI actions are expressed as agent “intents.”
and how model outputs should be structured for dynamic UI rendering

To make our ADK agent compatible with the AG-UI frontend, we need a way to intercept the agent’s internal events and convert them into AG-UI–compliant messages.

The AG-UI Python SDK provides a wrapper class that handles this translation for us. All we need to do is update our main application file to wrap the agent with this middleware and expose it through an AG-UI–compatible endpoint.

Below is the main.py file that initializes the agent runtime:

cd ./../
touch main.py

from ag_ui_adk import ADKAgent, add_adk_fastapi_endpoint
from fastapi import FastAPI

from agents.atlas_agent.agent import root_agent as atlas_root_agent

# -------------------------------------------------------------------
# Wrap the agent inside an ADKAgent middleware
# This provides sessions, user identity, in-memory services,
# and the unified ADK API that frontend UI components expect.
# -------------------------------------------------------------------
ag_atlas_agent = ADKAgent(
    adk_agent=atlas_root_agent, # The core ADK agent
    app_name="atlas_app", # App identifier
    user_id="demo_user", # Mock user ID (replace in production)
    session_timeout_seconds=3600, # Session expiration
    use_in_memory_services=True # Enables in-memory RAG + storage
)

# Create the FastAPI application
app = FastAPI(title="ADK Middleware Atlas Chat")

# add custom routes
@app.get("/health")
async def health_check():
    return {"status": "ok"}

# -------------------------------------------------------------------
# Register an ADK-compliant endpoint with FastAPI.
# This exposes the chat API at "/".
# Your frontend (Next.js + CopilotKit) will call this endpoint.
# -------------------------------------------------------------------
add_adk_fastapi_endpoint(app, ag_atlas_agent, path="/")

# -------------------------------------------------------------------
# Run the development server using Uvicorn
# Only executes when running `python main.py`
# -------------------------------------------------------------------
if __name__ == ' __main__':
    import uvicorn
    uvicorn.run(
        "main:app",
        host="localhost",
        port=8000,
        reload=True, # Auto-reload on code changes
        workers=1 # Single worker recommended for MCP tools
    )

2.3. Standardizing Frontend-to-Backend Communication with AG-UI

With the backend wired up, the next step is to connect our Next.js frontend to the agent so it can send messages and receive responses. To do this, we create a Next.js route /api/copilotkit and initialize the Copilot Runtime inside its POST handler.

This route is responsible for:

relaying requests from CopilotKit UI components
forwarding those requests to the ADK agent
returning streaming or structured responses back to the frontend

import {
  CopilotRuntime,
  ExperimentalEmptyAdapter,
  copilotRuntimeNextJSAppRouterEndpoint,
} from "@copilotkit/runtime";
import { HttpAgent } from "@ag-ui/client";
import { NextRequest } from "next/server";

// Create a service adapter for the CopilotKit runtime
const serviceAdapter = new ExperimentalEmptyAdapter();
// Create the main CopilotRuntime instance that manages communication between the frontend and backend agents
const runtime = new CopilotRuntime({
// Define the agents that will be available to the frontend
    agents: {
// Configure the ADK agent connection
    atlas_agent: new HttpAgent({
      // Specify the URL where the ADK agent is running
      url: "http://localhost:8000/",
    })
},
});
// Export the POST handler for the API route
export const POST = async (req: NextRequest) => {
// Create the request handler using CopilotKit's Next.js helper
const { handleRequest } = copilotRuntimeNextJSAppRouterEndpoint({
    runtime, // The CopilotRuntime instance we configured
    serviceAdapter, // The service adapter for agent coordination
    endpoint: "/api/copilotkit", // The endpoint path (matches this file's location)
});
return handleRequest(req);
};

2.4. Integrating CopilotKit

Now that communication between frontend and backend is established, we must configure the CopilotKit provider so all UI components can access the runtime.

The Provider manages:

agent sessions
streaming updates
contextual prompts
requests made by Copilot UI components

Wrap your entire application with it inside layout.tsx:

import type { Metadata } from "next";
import { Geist, Geist_Mono } from "next/font/google";
import "./globals.css";

import { CopilotKit } from "@copilotkit/react-core";
import "@copilotkit/react-ui/styles.css";

const geistSans = Geist({
  variable: "--font-geist-sans",
  subsets: ["latin"],
});

const geistMono = Geist_Mono({
  variable: "--font-geist-mono",
  subsets: ["latin"],
});

export const metadata: Metadata = {
  title: "Create Next App",
  description: "Generated by create next app",
};

export default function RootLayout({
  children,
}: Readonly<{
  children: React.ReactNode;
}>) {
  return (
    <html lang="en">
      <body
        className={`${geistSans.variable} ${geistMono.variable} antialiased`}
      >
        <CopilotKit runtimeUrl="/api/copilotkit"
                    agent="atlas_agent"
                    showDevConsole={false}>
            {children}
        </CopilotKit>
      </body>
    </html>
  );
}

Once wrapped, any page can use CopilotKit’s built-in components.

2.5. Adding Copilot UI Components

CopilotKit includes prebuilt UI components for interacting with agents, such as:

CopilotChat
CopilotPopup
CopilotSidebar

Here’s an example using CopilotChat:

"use client";

import {CopilotChat} from "@copilotkit/react-ui";

export default function HomePage() {
  return (
    <main className="h-screen w-screen">
        <CopilotChat
            className="h-full rounded-2xl max-w-6xl mx-auto"
            labels={{initial: "Hi, I'm an agent. Want to chat?"}}
            suggestions={[{
                title: "Weather in New York",
                message: "What's the weather like in New York?"
            }]}
        />
    </main>
  );
}

2.6. Giving the Agent UI Control: Frontend Tools

One of CopilotKit’s most powerful features is frontend tools — functions that let the LLM trigger actions directly in your UI. These tools enable the agent to do things like update component state, modify styles, or adjust the appearance of your app interface.

The useFrontendTool hook allows you to define actions that the Agents can invoke through a handler function. This is the main mechanism that gives your agent the ability to perform client-side operations, whether that means updating React state or triggering custom side effects.

Each frontend tool is defined using three core elements:

Name and description — tells the AI what the tool does and when to use it
Parameters — a schema describing the inputs the tool accepts
Handler function — executed on the client when the AI calls the tool

Optionally, you can also include a render function to display custom UI inside the chat, such as tool results, status information, or visual feedback.

"use client";

import {CopilotChat} from "@copilotkit/react-ui";
import {useState} from "react";
import {useFrontendTool} from "@copilotkit/react-core";

export default function HomePage() {
    const [background, setBackground] = useState<string>(
        "--copilot-kit-background-color"
    );

    /* --------------------------------------------------------------------------------------------
     * CHANGE BACKGROUND TOOL
     * This tool allows the LLM to set the chat background.
     * ------------------------------------------------------------------------------------------*/
    useFrontendTool({
        name: "change_background",
        description:
            "Change the chat's background using any CSS background value (color, gradient, etc.).",
        parameters: [
            {
                name: "background",
                type: "string",
                description: "CSS background definition (colors, gradients, etc).",
            },
        ],

        // The tool handler executes when the LLM calls this tool.
        handler: ({background}) => {
            setBackground(background);

            return {
                status: "success",
                message: `Background changed to ${background}`,
            };
        },
    });

  return (
    <main className="h-screen w-screen" style={{background}} >
        <CopilotChat
            className="h-full rounded-2xl max-w-6xl mx-auto"
            labels={{initial: "Hi, I'm an agent. Want to chat?"}}
            suggestions={[{
                title: "Weather in New York",
                message: "What's the weather like in New York?"
            },{
                title: "Change background",
                message: "Change the background to a right-to-left gradient from blue to green."
            }]}
        />
    </main>
  );
}

Frontend tools turn your app from a simple chat interface into an active, UI-aware assistant capable of controlling and updating interface components in real time. For example, if you refresh the page and then type, “Can you change the background of the app to red?” (See animation below), that request is sent to your ADK-AG-UI–powered backend.

Because the frontend tool definitions are synchronized with the agent’s context, the agent understands exactly which UI capabilities are available and how to invoke them. It responds by emitting the appropriate function-call event through ADK.

This function call is then routed back to the client, where the useFrontendTool handler intercepts the ADK event and executes the corresponding UI action. In this case, the handler updates a state variable, triggering React to re-render the component HomePage with the new background color.

By following this same pattern, you can define as many frontend tools as needed, enabling increasingly rich, interactive behaviors driven directly by agent intent.

Frontend tools are great, but what about when we want to render a custom component based on the output of a function called in the backend? In the previous section, when we defined our agent, we equipped it with two tools: get_weather and get_location. The question now becomes: how can we intercept the output of those tool calls on the client side so we can render our own custom UI components?

This is exactly what we are going to implement next. Fortunately, AG-UI makes this process straightforward through the useRenderToolCall hook, which lets us capture tool-call responses and render fully customized React components based on the tool’s output.

    /* --------------------------------------------------------------------------------------------
     * RENDER WEATHER TOOL CALL
     * This visually renders the result of the get_weather tool.
     * ------------------------------------------------------------------------------------------*/
    useRenderToolCall({
        name: "get_weather",
        description: "Get the current weather for a specified location.",
        available: "disabled", // Using MCP or manually invoking elsewhere
        parameters: [{name: "location", type: "string", required: true}],
        render: ({args, status, result}) => {
           // The result variable is the function response
            /* STATUS: inProgress --------------------------------------------------*/
            if (status === "inProgress") {
                return (
                    <div className="bg-[#667eea] text-white p-4 rounded-lg max-w-md">
                        <span className="animate-spin">⚙️ Retrieving weather...</span>
                    </div>
                );
            }

            /* STATUS: complete ----------------------------------------------------*/
            if (status === "complete" && result) {
            // Her we can render our custom component
                return (
                    <div className="bg-white p-4 rounded-lg shadow-md max-w-md">
                        <h2 className="text-xl font-bold mb-2">
                            Weather in {args.location}
                        </h2>
                        <p className="text-gray-700">
                            Temperature: {result.result.temperature}°C
                        </p>
                        <p className="text-gray-700">
                            Condition: {result.result.condition}
                        </p>
                    </div>
                )
            }
            return null;
        },
    });

We’ve included the complete source code for this tutorial — including the GoogleMap and WeatherCard components—to create a more engaging visualization for displaying the weather and location data, in the following GitHub repository: https://github.com/haruiz/atlas_app.git

To use the GoogleMap component, make sure to create an .env file in your Next.js application and define the following environment variable:

NEXT_PUBLIC_GOOGLE_MAPS_API_KEY=<YOUR GOOGLE MAPS API KEY>

/* --------------------------------------------------------------------------------------------
    * RENDER PLACE LOCATION TOOL CALL
    * This visually renders the result of the get_place_location tool.
    * ------------------------------------------------------------------------------------------
    */

    useRenderToolCall({
        name: "get_place_location",
        description: "get the latitude and longitude of a place given its name.",
        available: "disabled",
        parameters: [{name: "place_name", type: "string", required: true}],
        render: ({args, status, result}) => {
            if (status === "inProgress") {
                return (
                    <div className="bg-[#667eea] text-white p-4 rounded-lg max-w-md">
                        <span className="animate-spin">⚙️ Retrieving location...</span>
                    </div>
                );
            }
            if (status === "complete" && result) {
                const {result: coords} = result;
                return <GoogleMap lat={coords?.latitude} lng={coords?.longitude}/>;
            }
            return null;
        }
    })

    /* --------------------------------------------------------------------------------------------
     * RENDER WEATHER TOOL CALL
     * This visually renders the result of the get_weather tool.
     * ------------------------------------------------------------------------------------------*/
    useRenderToolCall({
        name: "get_weather",
        description: "Get the current weather for a specified location.",
        available: "disabled", // Using MCP or manually invoking elsewhere
        parameters: [{name: "location", type: "string", required: true}],
        render: ({args, status, result : toolResponse}) => {
            /* STATUS: inProgress --------------------------------------------------*/
            if (status === "inProgress") {
                return (
                    <div className="bg-[#667eea] text-white p-4 rounded-lg max-w-md">
                        <span className="animate-spin">⚙️ Retrieving weather...</span>
                    </div>
                );
            }

            /* STATUS: complete ----------------------------------------------------*/
            if (status === "complete" && toolResponse) {
                const weatherResult: WeatherToolResult | null = toolResponse?.result || null;
                console.log("Weather Result:", weatherResult);
                if (!weatherResult) {
                    return (
                        <div className="bg-red-300 text-red-900 p-4 rounded-lg max-w-md">
                            <strong>⚠️ Error:</strong> Unable to retrieve weather data. Please try again.
                        </div>
                    );
                }
                // Choose color based on weather conditions
                const themeColor = getThemeColor(weatherResult.conditions);
                return (
                    <WeatherCard
                        location={args.location}
                        themeColor={themeColor}
                        result={weatherResult}
                        status={status || "complete"}
                    />
                );
            }

            return null;
        },
    });

At this point, if we run our application, we should see the following:

By the end of this walkthrough, we have built a fully functional agentic application that goes beyond a simple chat interface. Using ADK to define and orchestrate our agents, and AG-UI to standardize how agent execution is surfaced in the frontend, we created an experience where agents can reason, call tools, stream progress, render custom UI components, and interact with users in real time.

We encourage you to explore ADK and AG-UI more deeply: experiment with additional agents, introduce new frontend tools, render richer UI components, and test human-in-the-loop workflows in your own applications. The real power of this stack emerges when agents stop being isolated responders and start becoming collaborative, UI-aware participants in your application.

In upcoming posts, we’ll dive deeper into advanced patterns such as multi-agent coordination, cross-agent routing, long-running tasks, and richer UI orchestration with AG-UI, such as Agent-to-Agent (A2A) integration. Stay tuned, and follow along as we continue exploring how to build robust, human-centered agentic systems.

Acknowledgements

This work was supported by Google ML Developer Programs and the Google Developers Program , which provided Google Cloud credits and high-quality technical resources #AISprint