nvidia-nat-adk 1.3.0a20250925__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- nat/meta/pypi.md +21 -0
- nat/plugins/adk/__init__.py +20 -0
- nat/plugins/adk/adk_callback_handler.py +296 -0
- nat/plugins/adk/llm.py +50 -0
- nat/plugins/adk/register.py +25 -0
- nat/plugins/adk/tool_wrapper.py +159 -0
- nvidia_nat_adk-1.3.0a20250925.dist-info/METADATA +36 -0
- nvidia_nat_adk-1.3.0a20250925.dist-info/RECORD +11 -0
- nvidia_nat_adk-1.3.0a20250925.dist-info/WHEEL +5 -0
- nvidia_nat_adk-1.3.0a20250925.dist-info/entry_points.txt +2 -0
- nvidia_nat_adk-1.3.0a20250925.dist-info/top_level.txt +1 -0
nat/meta/pypi.md
ADDED
@@ -0,0 +1,21 @@
|
|
1
|
+
<!--
|
2
|
+
SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
3
|
+
SPDX-License-Identifier: Apache-2.0
|
4
|
+
|
5
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
6
|
+
you may not use this file except in compliance with the License.
|
7
|
+
You may obtain a copy of the License at
|
8
|
+
|
9
|
+
http://www.apache.org/licenses/LICENSE-2.0
|
10
|
+
|
11
|
+
Unless required by applicable law or agreed to in writing, software
|
12
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
13
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
14
|
+
See the License for the specific language governing permissions and
|
15
|
+
limitations under the License.
|
16
|
+
-->
|
17
|
+
|
18
|
+

|
19
|
+
|
20
|
+
# NVIDIA NeMo Agent Toolkit — Google ADK Subpackage
|
21
|
+
Subpackage providing Google ADK integration for the NVIDIA NeMo Agent Toolkit.
|
@@ -0,0 +1,20 @@
|
|
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.
|
15
|
+
|
16
|
+
# pylint: disable=unused-import
|
17
|
+
# flake8: noqa
|
18
|
+
# isort:skip_file
|
19
|
+
|
20
|
+
# Import any providers which need to be automatically registered here
|
@@ -0,0 +1,296 @@
|
|
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.
|
15
|
+
|
16
|
+
import copy
|
17
|
+
import logging
|
18
|
+
import threading
|
19
|
+
import time
|
20
|
+
from collections.abc import Callable
|
21
|
+
from typing import Any
|
22
|
+
|
23
|
+
import litellm
|
24
|
+
|
25
|
+
from nat.builder.context import Context
|
26
|
+
from nat.builder.framework_enum import LLMFrameworkEnum
|
27
|
+
from nat.data_models.intermediate_step import IntermediateStepPayload
|
28
|
+
from nat.data_models.intermediate_step import IntermediateStepType
|
29
|
+
from nat.data_models.intermediate_step import StreamEventData
|
30
|
+
from nat.data_models.intermediate_step import TraceMetadata
|
31
|
+
from nat.data_models.intermediate_step import UsageInfo
|
32
|
+
from nat.profiler.callbacks.base_callback_class import BaseProfilerCallback
|
33
|
+
from nat.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
|
34
|
+
|
35
|
+
logger = logging.getLogger(__name__)
|
36
|
+
|
37
|
+
|
38
|
+
class ADKProfilerHandler(BaseProfilerCallback):
|
39
|
+
"""
|
40
|
+
A callback manager/handler for Google ADK that intercepts calls to:
|
41
|
+
- Tools
|
42
|
+
- LLMs
|
43
|
+
to collect usage statistics (tokens, inputs, outputs, time intervals, etc.)
|
44
|
+
and store them in NeMo Agent Toolkit's usage_stats queue for subsequent analysis.
|
45
|
+
"""
|
46
|
+
|
47
|
+
def __init__(self) -> None:
|
48
|
+
super().__init__()
|
49
|
+
self._lock = threading.Lock()
|
50
|
+
self.last_call_ts = time.time()
|
51
|
+
self.step_manager = Context.get().intermediate_step_manager
|
52
|
+
|
53
|
+
# Original references to Google ADK Tool and LLM methods (for uninstrumenting if needed)
|
54
|
+
self._original_tool_call = None
|
55
|
+
self._original_llm_call = None
|
56
|
+
self._instrumented = False
|
57
|
+
|
58
|
+
def instrument(self) -> None:
|
59
|
+
"""
|
60
|
+
Monkey-patch the relevant Google ADK methods with usage-stat collection logic.
|
61
|
+
Assumes the 'google-adk' library is installed.
|
62
|
+
"""
|
63
|
+
|
64
|
+
if getattr(self, "_instrumented", False):
|
65
|
+
logger.debug("ADKProfilerHandler already instrumented; skipping.")
|
66
|
+
return
|
67
|
+
try:
|
68
|
+
from google.adk.tools.function_tool import FunctionTool
|
69
|
+
except Exception as _e:
|
70
|
+
logger.exception("ADK import failed; skipping instrumentation")
|
71
|
+
return
|
72
|
+
|
73
|
+
# Save the originals
|
74
|
+
self._original_tool_call = getattr(FunctionTool, "run_async", None)
|
75
|
+
self._original_llm_call = getattr(litellm, "acompletion", None)
|
76
|
+
|
77
|
+
# Patch if available
|
78
|
+
if self._original_tool_call:
|
79
|
+
FunctionTool.run_async = self._tool_use_monkey_patch()
|
80
|
+
|
81
|
+
if self._original_llm_call:
|
82
|
+
litellm.acompletion = self._llm_call_monkey_patch()
|
83
|
+
|
84
|
+
logger.debug("ADKProfilerHandler instrumentation applied successfully.")
|
85
|
+
self._instrumented = True
|
86
|
+
|
87
|
+
def uninstrument(self) -> None:
|
88
|
+
""" Restore the original Google ADK methods.
|
89
|
+
Add an explicit unpatch to avoid side-effects across tests/process lifetime.
|
90
|
+
"""
|
91
|
+
try:
|
92
|
+
from google.adk.tools.function_tool import FunctionTool
|
93
|
+
if self._original_tool_call:
|
94
|
+
FunctionTool.run_async = self._original_tool_call
|
95
|
+
if self._original_llm_call:
|
96
|
+
litellm.acompletion = self._original_llm_call
|
97
|
+
logger.debug("ADKProfilerHandler uninstrumented successfully.")
|
98
|
+
except Exception as _e:
|
99
|
+
logger.exception("Failed to uninstrument ADKProfilerHandler")
|
100
|
+
|
101
|
+
def _tool_use_monkey_patch(self) -> Callable[..., Any]:
|
102
|
+
"""
|
103
|
+
Returns a function that wraps calls to BaseTool.run_async with usage-logging.
|
104
|
+
"""
|
105
|
+
original_func = self._original_tool_call
|
106
|
+
|
107
|
+
async def wrapped_tool_use(base_tool_instance, *args, **kwargs) -> Any:
|
108
|
+
"""
|
109
|
+
Replicates _tool_use_wrapper logic without wrapt: collects usage stats,
|
110
|
+
calls the original, and captures output stats.
|
111
|
+
|
112
|
+
Args:
|
113
|
+
base_tool_instance (FunctionTool): The instance of the tool being called.
|
114
|
+
*args: Positional arguments to the tool.
|
115
|
+
**kwargs: Keyword arguments to the tool.
|
116
|
+
|
117
|
+
Returns:
|
118
|
+
Any: The result of the tool execution.
|
119
|
+
"""
|
120
|
+
now = time.time()
|
121
|
+
tool_name = ""
|
122
|
+
|
123
|
+
try:
|
124
|
+
tool_name = base_tool_instance.name
|
125
|
+
except Exception as _e:
|
126
|
+
logger.exception("Error getting tool name")
|
127
|
+
tool_name = ""
|
128
|
+
|
129
|
+
try:
|
130
|
+
# Pre-call usage event - safely extract kwargs args if present
|
131
|
+
kwargs_args = (kwargs.get("args", {}) if isinstance(kwargs.get("args"), dict) else {})
|
132
|
+
stats = IntermediateStepPayload(
|
133
|
+
event_type=IntermediateStepType.TOOL_START,
|
134
|
+
framework=LLMFrameworkEnum.ADK,
|
135
|
+
name=tool_name,
|
136
|
+
data=StreamEventData(),
|
137
|
+
metadata=TraceMetadata(tool_inputs={
|
138
|
+
"args": args, "kwargs": dict(kwargs_args)
|
139
|
+
}),
|
140
|
+
usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
|
141
|
+
)
|
142
|
+
|
143
|
+
self.step_manager.push_intermediate_step(stats)
|
144
|
+
|
145
|
+
with self._lock:
|
146
|
+
self.last_call_ts = now
|
147
|
+
|
148
|
+
# Call the original _use(...)
|
149
|
+
result = await original_func(base_tool_instance, *args, **kwargs)
|
150
|
+
now = time.time()
|
151
|
+
# Post-call usage stats - safely extract kwargs args if present
|
152
|
+
kwargs_args = (kwargs.get("args", {}) if isinstance(kwargs.get("args"), dict) else {})
|
153
|
+
usage_stat = IntermediateStepPayload(
|
154
|
+
event_type=IntermediateStepType.TOOL_END,
|
155
|
+
span_event_timestamp=now,
|
156
|
+
framework=LLMFrameworkEnum.ADK,
|
157
|
+
name=tool_name,
|
158
|
+
data=StreamEventData(
|
159
|
+
input={
|
160
|
+
"args": args, "kwargs": dict(kwargs_args)
|
161
|
+
},
|
162
|
+
output=str(result),
|
163
|
+
),
|
164
|
+
metadata=TraceMetadata(tool_outputs={"result": str(result)}),
|
165
|
+
usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
|
166
|
+
)
|
167
|
+
|
168
|
+
self.step_manager.push_intermediate_step(usage_stat)
|
169
|
+
|
170
|
+
return result
|
171
|
+
|
172
|
+
except Exception as _e:
|
173
|
+
logger.exception("BaseTool error occured")
|
174
|
+
raise
|
175
|
+
|
176
|
+
return wrapped_tool_use
|
177
|
+
|
178
|
+
def _llm_call_monkey_patch(self) -> Callable[..., Any]:
|
179
|
+
"""
|
180
|
+
Returns a function that wraps calls to litellm.acompletion(...) with usage-logging.
|
181
|
+
|
182
|
+
Returns:
|
183
|
+
Callable[..., Any]: The wrapped function.
|
184
|
+
"""
|
185
|
+
original_func = self._original_llm_call
|
186
|
+
|
187
|
+
async def wrapped_llm_call(*args, **kwargs) -> Any:
|
188
|
+
"""
|
189
|
+
Replicates _llm_call_wrapper logic without wrapt: collects usage stats,
|
190
|
+
calls the original, and captures output stats.
|
191
|
+
|
192
|
+
Args:
|
193
|
+
*args: Positional arguments to the LLM call.
|
194
|
+
**kwargs: Keyword arguments to the LLM call.
|
195
|
+
|
196
|
+
Returns:
|
197
|
+
Any: The result of the LLM call.
|
198
|
+
"""
|
199
|
+
|
200
|
+
now = time.time()
|
201
|
+
with self._lock:
|
202
|
+
seconds_between_calls = int(now - self.last_call_ts)
|
203
|
+
model_name = kwargs.get("model")
|
204
|
+
if not model_name and args:
|
205
|
+
first = args[0]
|
206
|
+
if isinstance(first, str):
|
207
|
+
model_name = first
|
208
|
+
model_name = model_name or ""
|
209
|
+
|
210
|
+
model_input = ""
|
211
|
+
try:
|
212
|
+
for message in kwargs.get("messages", []):
|
213
|
+
content = message.get("content", "")
|
214
|
+
if isinstance(content, list):
|
215
|
+
for part in content:
|
216
|
+
if isinstance(part, dict):
|
217
|
+
model_input += str(part.get("text", "")) # text parts
|
218
|
+
else:
|
219
|
+
model_input += str(part)
|
220
|
+
else:
|
221
|
+
model_input += content or ""
|
222
|
+
except Exception as _e:
|
223
|
+
logger.exception("Error getting model input")
|
224
|
+
|
225
|
+
# Record the start event
|
226
|
+
input_stats = IntermediateStepPayload(
|
227
|
+
event_type=IntermediateStepType.LLM_START,
|
228
|
+
framework=LLMFrameworkEnum.ADK,
|
229
|
+
name=model_name,
|
230
|
+
data=StreamEventData(input=model_input),
|
231
|
+
metadata=TraceMetadata(chat_inputs=copy.deepcopy(kwargs.get("messages", []))),
|
232
|
+
usage_info=UsageInfo(
|
233
|
+
token_usage=TokenUsageBaseModel(),
|
234
|
+
num_llm_calls=1,
|
235
|
+
seconds_between_calls=seconds_between_calls,
|
236
|
+
),
|
237
|
+
)
|
238
|
+
|
239
|
+
self.step_manager.push_intermediate_step(input_stats)
|
240
|
+
|
241
|
+
# Call the original litellm.acompletion(...)
|
242
|
+
output = await original_func(*args, **kwargs)
|
243
|
+
|
244
|
+
model_output = ""
|
245
|
+
try:
|
246
|
+
for choice in output.choices:
|
247
|
+
msg = choice.message
|
248
|
+
model_output += msg.content or ""
|
249
|
+
except Exception as _e:
|
250
|
+
logger.exception("Error getting model output")
|
251
|
+
|
252
|
+
now = time.time()
|
253
|
+
# Record the end event
|
254
|
+
# Prepare safe metadata and usage
|
255
|
+
chat_resp: dict[str, Any] = {}
|
256
|
+
try:
|
257
|
+
if getattr(output, "choices", []):
|
258
|
+
first_choice = output.choices[0]
|
259
|
+
chat_resp = first_choice.model_dump() if hasattr(
|
260
|
+
first_choice, "model_dump") else getattr(first_choice, "__dict__", {}) or {}
|
261
|
+
except Exception as _e:
|
262
|
+
logger.exception("Error preparing chat_responses")
|
263
|
+
|
264
|
+
usage_payload: dict[str, Any] = {}
|
265
|
+
try:
|
266
|
+
usage_obj = getattr(output, "usage", None) or (getattr(output, "model_extra", {}) or {}).get("usage")
|
267
|
+
if usage_obj:
|
268
|
+
if hasattr(usage_obj, "model_dump"):
|
269
|
+
usage_payload = usage_obj.model_dump()
|
270
|
+
elif isinstance(usage_obj, dict):
|
271
|
+
usage_payload = usage_obj
|
272
|
+
except Exception as _e:
|
273
|
+
logger.exception("Error preparing token usage")
|
274
|
+
|
275
|
+
output_stats = IntermediateStepPayload(
|
276
|
+
event_type=IntermediateStepType.LLM_END,
|
277
|
+
span_event_timestamp=now,
|
278
|
+
framework=LLMFrameworkEnum.ADK,
|
279
|
+
name=model_name,
|
280
|
+
data=StreamEventData(input=model_input, output=model_output),
|
281
|
+
metadata=TraceMetadata(chat_responses=chat_resp),
|
282
|
+
usage_info=UsageInfo(
|
283
|
+
token_usage=TokenUsageBaseModel(**usage_payload),
|
284
|
+
num_llm_calls=1,
|
285
|
+
seconds_between_calls=seconds_between_calls,
|
286
|
+
),
|
287
|
+
)
|
288
|
+
|
289
|
+
self.step_manager.push_intermediate_step(output_stats)
|
290
|
+
|
291
|
+
with self._lock:
|
292
|
+
self.last_call_ts = now
|
293
|
+
|
294
|
+
return output
|
295
|
+
|
296
|
+
return wrapped_llm_call
|
nat/plugins/adk/llm.py
ADDED
@@ -0,0 +1,50 @@
|
|
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.
|
15
|
+
|
16
|
+
from collections.abc import AsyncIterator
|
17
|
+
|
18
|
+
from nat.builder.builder import Builder
|
19
|
+
from nat.builder.framework_enum import LLMFrameworkEnum
|
20
|
+
from nat.cli.register_workflow import register_llm_client
|
21
|
+
from nat.llm.litellm_llm import LiteLlmModelConfig
|
22
|
+
|
23
|
+
|
24
|
+
@register_llm_client(config_type=LiteLlmModelConfig, wrapper_type=LLMFrameworkEnum.ADK)
|
25
|
+
async def litellm_adk(
|
26
|
+
litellm_config: LiteLlmModelConfig,
|
27
|
+
_builder: Builder, # pylint: disable=W0613 (_builder not used)
|
28
|
+
) -> AsyncIterator["LiteLlm"]: # type: ignore # noqa: F821 (forward reference of LiteLlm)
|
29
|
+
"""Create and yield a Google ADK `LiteLlm` client from a NAT `LiteLlmModelConfig`.
|
30
|
+
|
31
|
+
Args:
|
32
|
+
litellm_config (LiteLlmModelConfig): The configuration for the LiteLlm model.
|
33
|
+
_builder (Builder): The NAT builder instance.
|
34
|
+
|
35
|
+
Yields:
|
36
|
+
AsyncIterator[LiteLlm]: An async iterator that yields a LiteLlm client.
|
37
|
+
"""
|
38
|
+
|
39
|
+
try:
|
40
|
+
from google.adk.models.lite_llm import LiteLlm
|
41
|
+
except ModuleNotFoundError as e:
|
42
|
+
raise ModuleNotFoundError("Google ADK not installed; pip install google-adk") from e
|
43
|
+
|
44
|
+
llm = LiteLlm(**litellm_config.model_dump(
|
45
|
+
exclude={"type", "max_retries"},
|
46
|
+
by_alias=True,
|
47
|
+
exclude_none=True,
|
48
|
+
))
|
49
|
+
|
50
|
+
yield llm
|
@@ -0,0 +1,25 @@
|
|
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.
|
15
|
+
|
16
|
+
# pylint: disable=unused-import
|
17
|
+
# flake8: noqa
|
18
|
+
# isort:skip_file
|
19
|
+
# Auto-registration module: importing submodules triggers provider registration via side effects.
|
20
|
+
"""Register ADK providers via import side-effects (loaded from the `nat_adk` entry point)."""
|
21
|
+
|
22
|
+
# Import any providers which need to be automatically registered here
|
23
|
+
|
24
|
+
from . import llm # pylint: disable=unused-import # imported for side effects (registration)
|
25
|
+
from . import tool_wrapper # pylint: disable=unused-import # imported for side effects (registration)
|
@@ -0,0 +1,159 @@
|
|
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.
|
15
|
+
"""Tool Wrapper file"""
|
16
|
+
import logging
|
17
|
+
import types
|
18
|
+
from collections.abc import AsyncIterator
|
19
|
+
from collections.abc import Callable
|
20
|
+
from dataclasses import is_dataclass
|
21
|
+
from typing import Any
|
22
|
+
from typing import get_args
|
23
|
+
from typing import get_origin
|
24
|
+
|
25
|
+
from pydantic import BaseModel
|
26
|
+
|
27
|
+
from nat.builder.builder import Builder
|
28
|
+
from nat.builder.framework_enum import LLMFrameworkEnum
|
29
|
+
from nat.builder.function import Function
|
30
|
+
from nat.cli.register_workflow import register_tool_wrapper
|
31
|
+
|
32
|
+
logger = logging.getLogger(__name__)
|
33
|
+
|
34
|
+
|
35
|
+
def resolve_type(t: Any) -> Any:
|
36
|
+
"""Return the non-None member of a Union/PEP 604 union;
|
37
|
+
otherwise return the type unchanged.
|
38
|
+
|
39
|
+
Args:
|
40
|
+
t (Any): The type to resolve.
|
41
|
+
|
42
|
+
Returns:
|
43
|
+
Any: The resolved type.
|
44
|
+
"""
|
45
|
+
origin = get_origin(t)
|
46
|
+
if origin is types.UnionType:
|
47
|
+
for arg in get_args(t):
|
48
|
+
if arg is not type(None):
|
49
|
+
return arg
|
50
|
+
return t
|
51
|
+
return t
|
52
|
+
|
53
|
+
|
54
|
+
@register_tool_wrapper(wrapper_type=LLMFrameworkEnum.ADK)
|
55
|
+
def google_adk_tool_wrapper(
|
56
|
+
name: str,
|
57
|
+
fn: Function,
|
58
|
+
_builder: Builder # pylint: disable=W0613
|
59
|
+
) -> Any: # Changed from Callable[..., Any] to Any to allow FunctionTool return
|
60
|
+
"""Wrap a NAT `Function` as a Google ADK `FunctionTool`.
|
61
|
+
|
62
|
+
Args:
|
63
|
+
name (str): The name of the tool.
|
64
|
+
fn (Function): The NAT `Function` to wrap.
|
65
|
+
_builder (Builder): The NAT `Builder` (not used).
|
66
|
+
Returns:
|
67
|
+
A Google ADK `FunctionTool` wrapping the NAT `Function`.
|
68
|
+
"""
|
69
|
+
import inspect
|
70
|
+
|
71
|
+
async def callable_ainvoke(*args: Any, **kwargs: Any) -> Any:
|
72
|
+
"""Async function to invoke the NAT function.
|
73
|
+
|
74
|
+
Args:
|
75
|
+
*args: Positional arguments to pass to the NAT function.
|
76
|
+
**kwargs: Keyword arguments to pass to the NAT function.
|
77
|
+
Returns:
|
78
|
+
Any: The result of invoking the NAT function.
|
79
|
+
"""
|
80
|
+
return await fn.acall_invoke(*args, **kwargs)
|
81
|
+
|
82
|
+
async def callable_astream(*args: Any, **kwargs: Any) -> AsyncIterator[Any]:
|
83
|
+
"""Async generator to stream results from the NAT function.
|
84
|
+
|
85
|
+
Args:
|
86
|
+
*args: Positional arguments to pass to the NAT function.
|
87
|
+
**kwargs: Keyword arguments to pass to the NAT function.
|
88
|
+
Yields:
|
89
|
+
Any: Streamed items from the NAT function.
|
90
|
+
"""
|
91
|
+
async for item in fn.acall_stream(*args, **kwargs):
|
92
|
+
yield item
|
93
|
+
|
94
|
+
def nat_function(
|
95
|
+
func: Callable[..., Any] | None = None,
|
96
|
+
*,
|
97
|
+
name: str = name,
|
98
|
+
description: str | None = fn.description,
|
99
|
+
input_schema: Any = fn.input_schema,
|
100
|
+
) -> Callable[..., Any]:
|
101
|
+
"""
|
102
|
+
Decorator to wrap a function as a NAT function.
|
103
|
+
|
104
|
+
Args:
|
105
|
+
func (Callable): The function to wrap.
|
106
|
+
name (str): The name of the function.
|
107
|
+
description (str): The description of the function.
|
108
|
+
input_schema (BaseModel): The Pydantic model defining the input schema.
|
109
|
+
|
110
|
+
Returns:
|
111
|
+
Callable[..., Any]: The wrapped function.
|
112
|
+
"""
|
113
|
+
if func is None:
|
114
|
+
raise ValueError("'func' must be provided.")
|
115
|
+
|
116
|
+
# If input_schema is a dataclass, convert it to a Pydantic model
|
117
|
+
if input_schema is not None and is_dataclass(input_schema):
|
118
|
+
input_schema = BaseModel.model_validate(input_schema)
|
119
|
+
|
120
|
+
def decorator(func_to_wrap: Callable[..., Any]) -> Callable[..., Any]:
|
121
|
+
"""
|
122
|
+
Decorator to set metadata on the function.
|
123
|
+
"""
|
124
|
+
# Set the function's metadata
|
125
|
+
if name is not None:
|
126
|
+
func_to_wrap.__name__ = name
|
127
|
+
if description is not None:
|
128
|
+
func_to_wrap.__doc__ = description
|
129
|
+
|
130
|
+
# Set signature only if input_schema is provided
|
131
|
+
params: list[inspect.Parameter] = []
|
132
|
+
if input_schema is not None:
|
133
|
+
annotations = getattr(input_schema, "__annotations__", {}) or {}
|
134
|
+
for param_name, param_annotation in annotations.items():
|
135
|
+
params.append(
|
136
|
+
inspect.Parameter(
|
137
|
+
param_name,
|
138
|
+
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
139
|
+
annotation=resolve_type(param_annotation),
|
140
|
+
))
|
141
|
+
setattr(func_to_wrap, "__signature__", inspect.Signature(parameters=params))
|
142
|
+
|
143
|
+
return func_to_wrap
|
144
|
+
|
145
|
+
# If func is None, return the decorator itself to be applied later
|
146
|
+
if func is None:
|
147
|
+
return decorator
|
148
|
+
# Otherwise, apply the decorator to the provided function
|
149
|
+
return decorator(func)
|
150
|
+
|
151
|
+
from google.adk.tools.function_tool import FunctionTool
|
152
|
+
|
153
|
+
if fn.has_streaming_output and not fn.has_single_output:
|
154
|
+
logger.debug("Creating streaming FunctionTool for: %s", name)
|
155
|
+
callable_tool = nat_function(func=callable_astream)
|
156
|
+
else:
|
157
|
+
logger.debug("Creating non-streaming FunctionTool for: %s", name)
|
158
|
+
callable_tool = nat_function(func=callable_ainvoke)
|
159
|
+
return FunctionTool(callable_tool)
|
@@ -0,0 +1,36 @@
|
|
1
|
+
Metadata-Version: 2.4
|
2
|
+
Name: nvidia-nat-adk
|
3
|
+
Version: 1.3.0a20250925
|
4
|
+
Summary: Subpackage for Google ADK integration in NeMo Agent Toolkit
|
5
|
+
Keywords: ai,rag,agents
|
6
|
+
Classifier: Programming Language :: Python
|
7
|
+
Classifier: Programming Language :: Python :: 3.11
|
8
|
+
Classifier: Programming Language :: Python :: 3.12
|
9
|
+
Classifier: Programming Language :: Python :: 3.13
|
10
|
+
Requires-Python: <3.14,>=3.11
|
11
|
+
Description-Content-Type: text/markdown
|
12
|
+
Requires-Dist: nvidia-nat==v1.3.0a20250925
|
13
|
+
Requires-Dist: google-adk~=1.6
|
14
|
+
Requires-Dist: litellm<1.74,>=1.63
|
15
|
+
|
16
|
+
<!--
|
17
|
+
SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
18
|
+
SPDX-License-Identifier: Apache-2.0
|
19
|
+
|
20
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
21
|
+
you may not use this file except in compliance with the License.
|
22
|
+
You may obtain a copy of the License at
|
23
|
+
|
24
|
+
http://www.apache.org/licenses/LICENSE-2.0
|
25
|
+
|
26
|
+
Unless required by applicable law or agreed to in writing, software
|
27
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
28
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
29
|
+
See the License for the specific language governing permissions and
|
30
|
+
limitations under the License.
|
31
|
+
-->
|
32
|
+
|
33
|
+

|
34
|
+
|
35
|
+
# NVIDIA NeMo Agent Toolkit — Google ADK Subpackage
|
36
|
+
Subpackage providing Google ADK integration for the NVIDIA NeMo Agent Toolkit.
|
@@ -0,0 +1,11 @@
|
|
1
|
+
nat/meta/pypi.md,sha256=AOSYTEk4JQQeaQQ1fA7v__j9-5rchBPvvSAICORgttE,953
|
2
|
+
nat/plugins/adk/__init__.py,sha256=bcb2ATld-lJ85wuXZWznP79-jfVZQp6SLwvSrt_vExU,817
|
3
|
+
nat/plugins/adk/adk_callback_handler.py,sha256=cZb2-CRKGKX9RMtveG5oYFQCDPMj43jxIdi3LgXEPYw,11803
|
4
|
+
nat/plugins/adk/llm.py,sha256=91zMW14t80tl-tMAmq-zzirPYBMnXaiv3dtCOLjMXr4,1936
|
5
|
+
nat/plugins/adk/register.py,sha256=-H9gEfqRRBUINJx1zvZiLc06oY2aXu-y1lcKxrF67Ts,1209
|
6
|
+
nat/plugins/adk/tool_wrapper.py,sha256=TroOo7tM6I3gOqMxOL7fQmPsyVKlntPPRNbYQMk6JHs,5779
|
7
|
+
nvidia_nat_adk-1.3.0a20250925.dist-info/METADATA,sha256=fBPVYAuCRt_iJLFdfaZdEvrSrA-AbhnzZamvjv5FvZk,1489
|
8
|
+
nvidia_nat_adk-1.3.0a20250925.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
9
|
+
nvidia_nat_adk-1.3.0a20250925.dist-info/entry_points.txt,sha256=K1K0clA1uqXk1Q0dh8ZaWNQyveaFx_XvgBWzCyGMYts,52
|
10
|
+
nvidia_nat_adk-1.3.0a20250925.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
|
11
|
+
nvidia_nat_adk-1.3.0a20250925.dist-info/RECORD,,
|
@@ -0,0 +1 @@
|
|
1
|
+
nat
|