Skip to content

Commit 3f6c917

Browse files
willkill07willkill07
authored
fix: improve Google ADK structure and fix callback handlers for tools and LLMs (#848)
### Google ADK Integration Improvements - fix LLM and TOOL call to ensure UUID matching - move agent code to `nvidia-nat-adk` from the example - generalize the ADK agent code - minor code clean up -- remove unnecessary classes - update example structure + content - add references to ADK in documentation - add missing `nvidia-nat-adk` reference in `nvidia-nat-all` - update test to work in the case of non-module import ### General Improvements - align/unify `litellm` and `openai` packages across all optional NAT packages Closes ## By Submitting this PR I confirm: - I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing.md). - We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. - Any contribution which contains commits that are not Signed-Off will not be accepted. - When the PR is ready for review, new or existing tests cover these changes. - When the PR is ready for review, the documentation is up to date with these changes. ## Summary by CodeRabbit * **New Features** * ADK-powered conversational agent with streaming responses and correlated step tracing (UUIDs). * New ADK install extras/package to simplify ADK integration. * **Refactor** * Streamlined LLM/tool configuration and registration; simplified time/weather tool behavior and removed legacy demo pathway. * **Documentation** * Updated demo README, quick-start, profiler, and feature docs to include ADK and clearer setup/output examples. Authors: - Will Killian (https://github.com/willkill07) Approvers: - David Gardner (https://github.com/dagardner-nv) URL: #848
1 parent ffdecac commit 3f6c917

File tree

30 files changed

+397
-483
lines changed

30 files changed

+397
-483
lines changed

‎docs/source/index.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ NeMo Agent toolkit was previously known as <!-- vale off -->AgentIQ<!-- vale on
3636

3737
## Key Features
3838

39-
- [**Framework Agnostic:**](./quick-start/installing.md#framework-integrations) NeMo Agent toolkit works side-by-side and around existing agentic frameworks, such as [LangChain](https://www.langchain.com/), [LlamaIndex](https://www.llamaindex.ai/), [CrewAI](https://www.crewai.com/), and [Microsoft Semantic Kernel](https://learn.microsoft.com/en-us/semantic-kernel/), as well as customer enterprise frameworks and simple Python agents. This allows you to use your current technology stack without replatforming. NeMo Agent toolkit complements any existing agentic framework or memory tool you're using and isn't tied to any specific agentic framework, long-term memory, or data source.
39+
- [**Framework Agnostic:**](./quick-start/installing.md#framework-integrations) NeMo Agent toolkit works side-by-side and around existing agentic frameworks, such as [LangChain](https://www.langchain.com/), [LlamaIndex](https://www.llamaindex.ai/), [CrewAI](https://www.crewai.com/), [Microsoft Semantic Kernel](https://learn.microsoft.com/en-us/semantic-kernel/), [Google ADK](https://github.com/google/adk-python), as well as customer enterprise frameworks and simple Python agents. This allows you to use your current technology stack without replatforming. NeMo Agent toolkit complements any existing agentic framework or memory tool you're using and isn't tied to any specific agentic framework, long-term memory, or data source.
4040

4141
- [**Reusability:**](./extend/sharing-components.md) Every agent, tool, and agentic workflow in this library exists as a function call that works together in complex software applications. The composability between these agents, tools, and workflows allows you to build once and reuse in different scenarios.
4242

‎docs/source/quick-start/installing.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ To keep the library lightweight, many of the first-party plugins supported by Ne
3434

3535
To install these first-party plugin libraries, you can use the full distribution name (for example, `nvidia-nat-langchain`) or use the `nvidia-nat[langchain]` extra distribution. The following extras are supported:
3636

37+
- `nvidia-nat[adk]` or `nvidia-nat-adk` - [Google ADK](https://github.com/google/adk-python)
3738
- `nvidia-nat[agno]` or `nvidia-nat-agno` - [Agno](https://agno.com/)
3839
- `nvidia-nat[crewai]` or `nvidia-nat-crewai` - [CrewAI](https://www.crewai.com/)
3940
- `nvidia-nat[data-flywheel]` or `nvidia-nat-data-flywheel` - [NeMo DataFlywheel](https://github.com/NVIDIA-AI-Blueprints/data-flywheel)

‎docs/source/workflows/observe/index.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ Each exporter can optionally include a processing pipeline that transforms, filt
146146
#### Integration Components
147147

148148
- **{py:class}`nat.profiler.decorators`**: Decorators that wrap workflow and LLM framework context managers to inject usage-collection callbacks.
149-
- **{py:class}`~nat.profiler.callbacks`**: Callback handlers that track usage statistics (tokens, time, inputs/outputs) and push them to the event stream. Supports LangChain/LangGraph, LLama Index, CrewAI, and Semantic Kernel frameworks.
149+
- **{py:class}`~nat.profiler.callbacks`**: Callback handlers that track usage statistics (tokens, time, inputs/outputs) and push them to the event stream. Supports LangChain/LangGraph, LLama Index, CrewAI, Semantic Kernel, and Google ADK frameworks.
150150

151151
### Registering a New Telemetry Provider as a Plugin
152152

‎docs/source/workflows/profiler.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ The NeMo Agent toolkit Profiler can be broken into the following components:
5050
### Profiler Decorators and Callbacks
5151
- `src/nat/profiler/decorators` directory defines decorators that can wrap each workflow or LLM framework context manager to inject usage-collection callbacks.
5252
- `src/nat/profiler/callbacks` directory implements callback handlers. These handlers track usage statistics (tokens, time, inputs/outputs) and push them to the NeMo Agent toolkit usage stats queue. We currently support callback handlers for LangChain/LangGraph,
53-
LlamaIndex, CrewAI, and Semantic Kernel.
53+
LlamaIndex, CrewAI, Google ADK, and Semantic Kernel.
5454

5555
### Profiler Runner
5656

‎examples/frameworks/adk_demo/README.md‎

Lines changed: 26 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -25,31 +25,25 @@ If you have not already done so, follow the instructions in the [Install Guide](
2525

2626
### Install this Workflow
2727

28-
From the root directory of the NAT library, run the following commands:
28+
From the root directory of the NAT library, run the following command:
2929

3030
```bash
31-
uv pip install -e '.[adk]' --prerelease=allow
3231
uv pip install -e examples/frameworks/adk_demo
3332
```
3433

3534
### Set up API keys
3635

37-
LiteLLM routes to many providers. For OpenAI:
36+
For this example, an OpenAI API key is required. You can set it as follows:
3837
```bash
3938
export OPENAI_API_KEY="<your_openai_key>"
4039
# Optional (defaults to https://api.openai.com/v1 if unset)
41-
export OPENAI_API_BASE="${OPENAI_API_BASE:-https://api.openai.com/v1}"
40+
export OPENAI_API_BASE="<your_openai_base_url>"
4241
```
43-
For Azure OpenAI, set:
44-
```bash
45-
export OPENAI_API_KEY="<your_azure_openai_key>"
46-
export OPENAI_API_BASE="https://<your-azure-endpoint>/openai"
47-
```
48-
You can find LLM provider specific instructions in the LiteLLM documentation. Please set the appropriate environment variables.
4942

50-
### Run the Workflow
43+
Google ADK support within NeMo Agent toolkit currently only supports OpenAI and Azure OpenAI models for tool calling.
44+
45+
### Set up the MCP Server
5146

52-
#### Set up the MCP server
5347
This example also demonstrates how NAT can interact with MCP servers on behalf of ADK.
5448

5549
First run the MCP server with this command.
@@ -58,31 +52,36 @@ First run the MCP server with this command.
5852
nat mcp serve --config_file examples/frameworks/adk_demo/configs/config.yml --host 0.0.0.0 --port 9901 --name "My MCP Server"
5953
```
6054

55+
## Run the Workflow
56+
6157
Then run the workflow with the NAT CLI
6258

6359
```bash
6460
nat run --config_file examples/frameworks/adk_demo/configs/config.yml --input "What is the weather and time in New York today?"
6561
```
6662

67-
### Expected output
63+
### Expected Output
6864

6965
```console
70-
(.venv)
71-
[12:44] BASH_$
72-
> nat run --config_file examples/frameworks/adk_demo/configs/config.yml --input "What is the weather and time in New York today?"
73-
12:44:06 - LiteLLM:INFO: cost_calculator.py:588 - selected model name for cost calculation: openai/gpt-4.1-2025-04-14
74-
2025.09.19_12:44:06 || INFO || LiteLLM:588 :: selected model name for cost calculation: openai/gpt-4.1-2025-04-14
75-
selected model name for cost calculation: openai/gpt-4.1-2025-04-14
76-
2025.09.19_12:44:06 || INFO || nat.front_ends.console.console_front_end_plugin:96 ::
77-
--------------------------------------------------
78-
Workflow Result:
79-
['Today in New York:\n- The weather is sunny with a temperature of 25°C (77°F).\n- The current time is 3:44 PM (EDT).']
80-
--------------------------------------------------
66+
<snipped for brevity>
67+
68+
Configuration Summary:
69+
--------------------
70+
Workflow Type: adk
71+
Number of Functions: 2
72+
Number of Function Groups: 0
73+
Number of LLMs: 1
74+
Number of Embedders: 0
75+
Number of Memory: 0
76+
Number of Object Stores: 0
77+
Number of Retrievers: 0
78+
Number of TTC Strategies: 0
79+
Number of Authentication Providers: 0
80+
81+
<snipped for brevity>
8182

8283
--------------------------------------------------
8384
Workflow Result:
84-
['Today in New York:\n- The weather is sunny with a temperature of 25°C (77°F).\n- The current time is 3:44 PM (EDT).']
85+
['Here’s the latest for New York:\n- Weather: Sunny, around 25°C (77°F)\n- Time: 2025-09-25 12:27:26 EDT (UTC-4)']
8586
--------------------------------------------------
86-
87-
(.venv)
8887
```

‎examples/frameworks/adk_demo/configs/config.yml‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -36,15 +36,14 @@ functions:
3636
url: "http://localhost:9901/mcp"
3737

3838
llms:
39-
litellm_llm:
40-
_type: litellm
41-
model: gpt-5
42-
api_key: ${OPENAI_API_KEY}
43-
base_url: ${OPENAI_API_BASE:-https://api.openai.com/v1}
44-
temperature: 0.0
39+
llm:
40+
_type: openai
41+
model_name: gpt-5
4542

4643
workflow:
4744
_type: adk
45+
name: adk_demo
46+
llm: llm
4847
tool_names: [weather_update_tool, get_city_time_tool]
49-
llm_name: litellm_llm
5048
description: "To get the current weather and time in a specific city"
49+
prompt: "You are a helpful agent who can answer user questions about weather in a city. You also have a sub-agent that can answer questions about the time in a city."

‎examples/frameworks/adk_demo/configs/eval_config.yml‎

Lines changed: 26 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -13,18 +13,35 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515
# path-check-skip-file
16-
general:
17-
use_uvloop: true
18-
front_end:
19-
_type: fastapi
20-
cors:
21-
allow_origins: ["*"]
22-
step_adaptor:
23-
mode: default
16+
17+
functions:
18+
weather_update_tool:
19+
_type: nat_adk_demo/weather_update
20+
description: "Get the current weather for a specified city"
21+
get_city_time_tool:
22+
_type: nat_adk_demo/get_city_time_tool
23+
transport: "http"
24+
mcp_tool_name: get_city_time
25+
url: "http://localhost:9901/mcp"
26+
27+
llms:
28+
llm:
29+
_type: openai
30+
model_name: gpt-5
31+
32+
workflow:
33+
_type: adk
34+
name: adk_demo
35+
llm: llm
36+
tool_names: [weather_update_tool, get_city_time_tool]
37+
description: "To get the current weather and time in a specific city"
38+
prompt: "You are a helpful agent who can answer user questions about weather in a city. You also have a sub-agent that can answer questions about the time in a city."
2439

2540
eval:
2641
general:
27-
output_dir: nat_adk_demo/eval_output
42+
output:
43+
dir: ./.tmp/nat/examples/frameworks/adk_demo/eval/
44+
cleanup: true
2845
dataset:
2946
_type: json
3047
file_path: examples/frameworks/adk_demo/data/eval.json
@@ -41,25 +58,3 @@ eval:
4158
concurrency_spike_analysis:
4259
enable: true
4360
spike_threshold: 7
44-
45-
functions:
46-
weather_update_tool:
47-
_type: nat_adk_demo/weather_update
48-
description: "Get the current weather for a specified city"
49-
get_city_time_tool:
50-
_type: nat_adk_demo/get_city_time_tool
51-
transport: "http"
52-
mcp_tool_name: get_city_time
53-
url: "http://localhost:9901/mcp"
54-
55-
llms:
56-
litellm_llm:
57-
_type: nat_test_llm
58-
response_seq: ["The weather in New York is sunny. The current time is 12:00:00 EDT-0400."]
59-
delay_ms: 0
60-
61-
workflow:
62-
_type: adk
63-
tool_names: [weather_update_tool, get_city_time_tool]
64-
llm_name: litellm_llm
65-
description: "To get the current weather and time in a specific city"

‎examples/frameworks/adk_demo/data/eval.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
version https://git-lfs.github.com/spec/v1
2-
oid sha256:a3366f6c9ce43f2f1355c9ac37e2c20f1ca04cb9e0dedcef48b776ca7d47fb8f
3-
size 885
2+
oid sha256:e7e0a3ad0f16adbcf51b068d95b3e9c3a3d420da7d4b44724cd3d68e74380139
3+
size 762

‎examples/frameworks/adk_demo/src/nat_adk_demo/__init__.py‎

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +0,0 @@
1-
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2-
# SPDX-License-Identifier: Apache-2.0
3-
#
4-
# Licensed under the Apache License, Version 2.0 (the "License");
5-
# you may not use this file except in compliance with the License.
6-
# You may obtain a copy of the License at
7-
#
8-
# http://www.apache.org/licenses/LICENSE-2.0
9-
#
10-
# Unless required by applicable law or agreed to in writing, software
11-
# distributed under the License is distributed on an "AS IS" BASIS,
12-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13-
# See the License for the specific language governing permissions and
14-
# limitations under the License.

‎examples/frameworks/adk_demo/src/nat_adk_demo/nat_time_mcp_tool.py‎

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -32,33 +32,29 @@ class TimeMCPToolConfig(FunctionBaseConfig, name="get_city_time_tool"):
3232

3333
@register_function(config_type=TimeMCPToolConfig)
3434
async def get_city_time(_config: TimeMCPToolConfig, _builder: Builder) -> AsyncIterator[FunctionInfo]:
35-
"""Register the get_city_time(city: str) -> str MCP tool.
35+
"""
36+
Register a get_city_time(city: str) -> str tool for ADK.
3637
37-
Yields:
38-
FunctionInfo: Descriptor for an async function `_get_city_time(city: str) -> str`.
38+
Args:
39+
_config (TimeMCPToolConfig): The configuration for the get_city_time tool.
40+
_builder (Builder): The NAT builder instance.
3941
"""
4042

4143
async def _get_city_time(city: str) -> str:
42-
"""Get the time in a specified city.
43-
Currently supports New York.
44+
"""
45+
Get the time in a specified city.
46+
4447
Args:
4548
city (str): The name of the city.
4649
4750
Returns:
4851
str: The current time in the specified city or an error message if the city is not recognized.
4952
"""
5053

51-
if city.strip().casefold() in {"new york", "new york city", "nyc"}:
52-
tz_identifier = "America/New_York"
53-
else:
54+
if city.strip().casefold() not in {"new york", "new york city", "nyc"}:
5455
return f"Sorry, I don't have timezone information for {city}."
5556

56-
tz = ZoneInfo(tz_identifier)
57-
now = datetime.datetime.now(tz)
58-
report = f'The current time in {city} is {now.strftime("%Y-%m-%d %H:%M:%S %Z%z")}'
59-
return report
57+
now = datetime.datetime.now(ZoneInfo("America/New_York"))
58+
return f'The current time in {city} is {now.strftime("%Y-%m-%d %H:%M:%S %Z%z")}'
6059

61-
# Create a Generic NAT tool that can be used with any supported LLM framework
62-
yield FunctionInfo.from_fn(_get_city_time,
63-
description=("This tool provides the current time in a specified city. "
64-
"It takes a city name as input and returns the current time in that city."))
60+
yield FunctionInfo.from_fn(_get_city_time, description=_get_city_time.__doc__)

0 commit comments

Comments
 (0)