nvidia-nat-adk 1.3.0a20250925__py3-none-any.whl → 1.3.0a20250928__py3-none-any.whl

nat/plugins/adk/__init__.py

@@ -12,9 +12,3 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-# pylint: disable=unused-import
-# flake8: noqa
-# isort:skip_file
-
-# Import any providers which need to be automatically registered here

nat/plugins/adk/adk_callback_handler.py

@@ -20,8 +20,6 @@ import time
 from collections.abc import Callable
 from typing import Any
 
-import litellm
-
 from nat.builder.context import Context
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.data_models.intermediate_step import IntermediateStepPayload
@@ -60,6 +58,7 @@ class ADKProfilerHandler(BaseProfilerCallback):
         Monkey-patch the relevant Google ADK methods with usage-stat collection logic.
         Assumes the 'google-adk' library is installed.
         """
+        import litellm
 
         if getattr(self, "_instrumented", False):
             logger.debug("ADKProfilerHandler already instrumented; skipping.")
@@ -89,6 +88,7 @@ class ADKProfilerHandler(BaseProfilerCallback):
         Add an explicit unpatch to avoid side-effects across tests/process lifetime.
         """
         try:
+            import litellm
             from google.adk.tools.function_tool import FunctionTool
             if self._original_tool_call:
                 FunctionTool.run_async = self._original_tool_call
@@ -140,12 +140,17 @@ class ADKProfilerHandler(BaseProfilerCallback):
                     usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
                 )
 
+                # Store the UUID to ensure the END event uses the same ID
+                step_uuid = stats.UUID
                 self.step_manager.push_intermediate_step(stats)
 
                 with self._lock:
                     self.last_call_ts = now
 
                 # Call the original _use(...)
+                if original_func is None:
+                    raise RuntimeError(
+                        "Original tool function is None - instrumentation may not have been set up correctly")
                 result = await original_func(base_tool_instance, *args, **kwargs)
                 now = time.time()
                 # Post-call usage stats - safely extract kwargs args if present
@@ -163,6 +168,7 @@ class ADKProfilerHandler(BaseProfilerCallback):
                     ),
                     metadata=TraceMetadata(tool_outputs={"result": str(result)}),
                     usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
+                    UUID=step_uuid,  # Use the same UUID as the START event
                 )
 
                 self.step_manager.push_intermediate_step(usage_stat)
@@ -236,9 +242,13 @@ class ADKProfilerHandler(BaseProfilerCallback):
                 ),
             )
 
+            # Store the UUID to ensure the END event uses the same ID
+            step_uuid = input_stats.UUID
             self.step_manager.push_intermediate_step(input_stats)
 
             # Call the original litellm.acompletion(...)
+            if original_func is None:
+                raise RuntimeError("Original LLM function is None - instrumentation may not have been set up correctly")
             output = await original_func(*args, **kwargs)
 
             model_output = ""
@@ -284,6 +294,7 @@ class ADKProfilerHandler(BaseProfilerCallback):
                     num_llm_calls=1,
                     seconds_between_calls=seconds_between_calls,
                 ),
+                UUID=step_uuid,  # Use the same UUID as the START event
             )
 
             self.step_manager.push_intermediate_step(output_stats)

nat/plugins/adk/agent.py

@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ADKFunctionConfig(FunctionBaseConfig, name="adk"):
+    """Configuration for ADK demo function."""
+    name: str = Field(default="nat-adk-agent")
+    description: str
+    prompt: str
+    llm: LLMRef
+    tool_names: list[str] = Field(default_factory=list)
+    workflow_alias: str = Field(default="adk_agent")
+    user_id: str = Field(default="nat")
+
+
+@register_function(config_type=ADKFunctionConfig, framework_wrappers=[LLMFrameworkEnum.ADK])
+async def adk_agent(config: ADKFunctionConfig, builder: Builder):
+    """An example function that demonstrates how to use the Google ADK framework with NAT.
+
+    Args:
+        config (ADKFunctionConfig): The configuration for the ADK agent function.
+        builder (Builder): The NAT builder instance.
+    """
+    import logging
+
+    from google.adk import Runner
+    from google.adk.agents import Agent
+    from google.adk.artifacts import InMemoryArtifactService
+    from google.adk.sessions import InMemorySessionService
+    from google.genai import types
+
+    logging.getLogger("LiteLLM").setLevel(logging.WARNING)
+
+    model = await builder.get_llm(config.llm, wrapper_type=LLMFrameworkEnum.ADK)
+    tools = await builder.get_tools(config.tool_names, wrapper_type=LLMFrameworkEnum.ADK)
+
+    agent = Agent(
+        name=config.name,
+        model=model,
+        description=config.description,
+        instruction=config.prompt,
+        tools=tools,
+    )
+
+    # Initialize the Runner with the agent and services
+    session_service = InMemorySessionService()
+    artifact_service = InMemoryArtifactService()
+    runner = Runner(app_name=config.name,
+                    agent=agent,
+                    artifact_service=artifact_service,
+                    session_service=session_service)
+    session = await session_service.create_session(app_name=config.name, user_id=config.user_id)
+
+    async def _response_fn(input_message: str) -> str:
+        """Wrapper for response fn
+
+        Args:
+            input_message (str): The input message from the user.
+        Returns:
+            str : The response from the agent.
+        """
+
+        async def run_prompt(new_message: str) -> str:
+            """Run prompt through the agent.
+
+            Args:
+                new_message (str): The input message from the user.
+            Returns:
+                str: The response from the agent.
+            """
+            content = types.Content(role="user", parts=[types.Part.from_text(text=new_message)])
+            text_buf: list[str] = []
+            async for event in runner.run_async(user_id=config.user_id, session_id=session.id, new_message=content):
+                if event.content is None:
+                    continue
+                if event.content.parts is None:
+                    continue
+                text_buf.extend(part.text for part in event.content.parts if part.text is not None)
+            return "".join(text_buf) if text_buf else ""
+
+        return await run_prompt(input_message)
+
+    try:
+        yield FunctionInfo.create(single_fn=_response_fn, description=config.description)
+    except GeneratorExit:
+        logger.debug("Exited early!", exc_info=True)
+    finally:
+        logger.debug("Cleaning up")

nat/plugins/adk/llm.py

@@ -13,38 +13,96 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections.abc import AsyncIterator
+import logging
+import os
 
 from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.register_workflow import register_llm_client
+from nat.llm.azure_openai_llm import AzureOpenAIModelConfig
 from nat.llm.litellm_llm import LiteLlmModelConfig
+from nat.llm.nim_llm import NIMModelConfig
+from nat.llm.openai_llm import OpenAIModelConfig
 
+logger = logging.getLogger(__name__)
 
-@register_llm_client(config_type=LiteLlmModelConfig, wrapper_type=LLMFrameworkEnum.ADK)
-async def litellm_adk(
-    litellm_config: LiteLlmModelConfig,
-    _builder: Builder,  # pylint: disable=W0613 (_builder not used)
-) -> AsyncIterator["LiteLlm"]:  # type: ignore # noqa: F821 (forward reference of LiteLlm)
-    """Create and yield a Google ADK `LiteLlm` client from a NAT `LiteLlmModelConfig`.
+
+@register_llm_client(config_type=AzureOpenAIModelConfig, wrapper_type=LLMFrameworkEnum.ADK)
+async def azure_openai_adk(config: AzureOpenAIModelConfig, _builder: Builder):
+    """Create and yield a Google ADK `AzureOpenAI` client from a NAT `AzureOpenAIModelConfig`.
 
     Args:
-        litellm_config (LiteLlmModelConfig): The configuration for the LiteLlm model.
+        config (AzureOpenAIModelConfig): The configuration for the AzureOpenAI model.
         _builder (Builder): The NAT builder instance.
-
-    Yields:
-        AsyncIterator[LiteLlm]: An async iterator that yields a LiteLlm client.
     """
+    from google.adk.models.lite_llm import LiteLlm
+
+    config_dict = config.model_dump(
+        exclude={"type", "max_retries", "thinking", "azure_endpoint", "azure_deployment", "model_name", "model"},
+        by_alias=True,
+        exclude_none=True,
+    )
+    if config.azure_endpoint:
+        config_dict["api_base"] = config.azure_endpoint
 
-    try:
-        from google.adk.models.lite_llm import LiteLlm
-    except ModuleNotFoundError as e:
-        raise ModuleNotFoundError("Google ADK not installed; pip install google-adk") from e
+    yield LiteLlm(f"azure/{config.azure_deployment}", **config_dict)
 
-    llm = LiteLlm(**litellm_config.model_dump(
-        exclude={"type", "max_retries"},
+
+@register_llm_client(config_type=LiteLlmModelConfig, wrapper_type=LLMFrameworkEnum.ADK)
+async def litellm_adk(litellm_config: LiteLlmModelConfig, _builder: Builder):
+    from google.adk.models.lite_llm import LiteLlm
+    yield LiteLlm(**litellm_config.model_dump(
+        exclude={"type", "max_retries", "thinking"},
         by_alias=True,
         exclude_none=True,
     ))
 
-    yield llm
+
+@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.ADK)
+async def nim_adk(config: NIMModelConfig, _builder: Builder):
+    """Create and yield a Google ADK `NIM` client from a NAT `NIMModelConfig`.
+
+    Args:
+        config (NIMModelConfig): The configuration for the NIM model.
+        _builder (Builder): The NAT builder instance.
+    """
+    import litellm
+    from google.adk.models.lite_llm import LiteLlm
+
+    logger.warning("NIMs do not currently support tools with ADK. Tools will be ignored.")
+    litellm.add_function_to_prompt = True
+    litellm.drop_params = True
+
+    if (api_key := os.getenv("NVIDIA_API_KEY", None)) is not None:
+        os.environ["NVIDIA_NIM_API_KEY"] = api_key
+
+    config_dict = config.model_dump(
+        exclude={"type", "max_retries", "thinking", "model_name", "model", "base_url"},
+        by_alias=True,
+        exclude_none=True,
+    )
+    if config.base_url:
+        config_dict["api_base"] = config.base_url
+
+    yield LiteLlm(f"nvidia_nim/{config.model_name}", **config_dict)
+
+
+@register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.ADK)
+async def openai_adk(config: OpenAIModelConfig, _builder: Builder):
+    """Create and yield a Google ADK `OpenAI` client from a NAT `OpenAIModelConfig`.
+
+    Args:
+        config (OpenAIModelConfig): The configuration for the OpenAI model.
+        _builder (Builder): The NAT builder instance.
+    """
+    from google.adk.models.lite_llm import LiteLlm
+
+    config_dict = config.model_dump(
+        exclude={"type", "max_retries", "thinking", "model_name", "model", "base_url"},
+        by_alias=True,
+        exclude_none=True,
+    )
+    if config.base_url:
+        config_dict["api_base"] = config.base_url
+
+    yield LiteLlm(config.model_name, **config_dict)

nat/plugins/adk/register.py

@@ -15,11 +15,8 @@
 
 # pylint: disable=unused-import
 # flake8: noqa
-# isort:skip_file
-# Auto-registration module: importing submodules triggers provider registration via side effects.
-"""Register ADK providers via import side-effects (loaded from the `nat_adk` entry point)."""
 
-# Import any providers which need to be automatically registered here
-
-from . import llm  # pylint: disable=unused-import  # imported for side effects (registration)
-from . import tool_wrapper  # pylint: disable=unused-import  # imported for side effects (registration)
+from . import adk_callback_handler
+from . import agent
+from . import llm
+from . import tool_wrapper

{nvidia_nat_adk-1.3.0a20250925.dist-info → nvidia_nat_adk-1.3.0a20250928.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-adk
-Version: 1.3.0a20250925
+Version: 1.3.0a20250928
 Summary: Subpackage for Google ADK integration in NeMo Agent Toolkit
 Keywords: ai,rag,agents
 Classifier: Programming Language :: Python
@@ -9,9 +9,8 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: nvidia-nat==v1.3.0a20250925
-Requires-Dist: google-adk~=1.6
-Requires-Dist: litellm<1.74,>=1.63
+Requires-Dist: nvidia-nat[litellm]==v1.3.0a20250928
+Requires-Dist: google-adk~=1.14.1
 
 <!--
 SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

nvidia_nat_adk-1.3.0a20250928.dist-info/RECORD

@@ -0,0 +1,12 @@
+nat/meta/pypi.md,sha256=AOSYTEk4JQQeaQQ1fA7v__j9-5rchBPvvSAICORgttE,953
+nat/plugins/adk/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
+nat/plugins/adk/adk_callback_handler.py,sha256=mJIeyREdqlDCwIuvudT5FaEyMqioISqjBRPFUGJoLGA,12561
+nat/plugins/adk/agent.py,sha256=9EpPSocb9OdPX9p5I7nrfe11L1ldhdOk7v09NMVcHNw,4270
+nat/plugins/adk/llm.py,sha256=glwlO9UQhfGagVDNbtw-bNQ7FVdJfs31v6rqd-MBJLY,4120
+nat/plugins/adk/register.py,sha256=9EEYSA9CFkPiEwVX-FB5Ho16wzaFsyO0hiiOUJFmaiM,829
+nat/plugins/adk/tool_wrapper.py,sha256=TroOo7tM6I3gOqMxOL7fQmPsyVKlntPPRNbYQMk6JHs,5779
+nvidia_nat_adk-1.3.0a20250928.dist-info/METADATA,sha256=oCtgyABZx7WyXDap2J04RZ4c6-RRYJ4Q1P02KezLVlQ,1466
+nvidia_nat_adk-1.3.0a20250928.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_adk-1.3.0a20250928.dist-info/entry_points.txt,sha256=K1K0clA1uqXk1Q0dh8ZaWNQyveaFx_XvgBWzCyGMYts,52
+nvidia_nat_adk-1.3.0a20250928.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_adk-1.3.0a20250928.dist-info/RECORD,,

nvidia_nat_adk-1.3.0a20250925.dist-info/RECORD

@@ -1,11 +0,0 @@
-nat/meta/pypi.md,sha256=AOSYTEk4JQQeaQQ1fA7v__j9-5rchBPvvSAICORgttE,953
-nat/plugins/adk/__init__.py,sha256=bcb2ATld-lJ85wuXZWznP79-jfVZQp6SLwvSrt_vExU,817
-nat/plugins/adk/adk_callback_handler.py,sha256=cZb2-CRKGKX9RMtveG5oYFQCDPMj43jxIdi3LgXEPYw,11803
-nat/plugins/adk/llm.py,sha256=91zMW14t80tl-tMAmq-zzirPYBMnXaiv3dtCOLjMXr4,1936
-nat/plugins/adk/register.py,sha256=-H9gEfqRRBUINJx1zvZiLc06oY2aXu-y1lcKxrF67Ts,1209
-nat/plugins/adk/tool_wrapper.py,sha256=TroOo7tM6I3gOqMxOL7fQmPsyVKlntPPRNbYQMk6JHs,5779
-nvidia_nat_adk-1.3.0a20250925.dist-info/METADATA,sha256=fBPVYAuCRt_iJLFdfaZdEvrSrA-AbhnzZamvjv5FvZk,1489
-nvidia_nat_adk-1.3.0a20250925.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_adk-1.3.0a20250925.dist-info/entry_points.txt,sha256=K1K0clA1uqXk1Q0dh8ZaWNQyveaFx_XvgBWzCyGMYts,52
-nvidia_nat_adk-1.3.0a20250925.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_adk-1.3.0a20250925.dist-info/RECORD,,