nvidia-nat-agno 1.3.dev0__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/plugins/agno/llm.py

@@ -13,75 +13,89 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
+from typing import TypeVar
 
 from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.register_workflow import register_llm_client
+from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.retry_mixin import RetryMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
 from nat.llm.nim_llm import NIMModelConfig
 from nat.llm.openai_llm import OpenAIModelConfig
+from nat.llm.utils.thinking import BaseThinkingInjector
+from nat.llm.utils.thinking import FunctionArgumentWrapper
+from nat.llm.utils.thinking import patch_with_thinking
 from nat.utils.exception_handlers.automatic_retries import patch_with_retry
+from nat.utils.type_utils import override
 
+ModelType = TypeVar("ModelType")
 
-@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.AGNO)
-async def nim_agno(llm_config: NIMModelConfig, builder: Builder):
-
-    from agno.models.nvidia import Nvidia
-
-    config_obj = {
-        **llm_config.model_dump(exclude={"type", "model_name"}, by_alias=True),
-        "id": f"{llm_config.model_name}",
-    }
-
-    # Because Agno uses a different environment variable for the API key, we need to set it here manually
-    if ("api_key" not in config_obj or config_obj["api_key"] is None):
 
-        if ("NVIDIA_API_KEY" in os.environ):
-            # Dont need to do anything. User has already set the correct key
-            pass
-        else:
-            nvidai_api_key = os.getenv("NVIDIA_API_KEY")
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
 
-            if (nvidai_api_key is not None):
-                # Transfer the key to the correct environment variable
-                os.environ["NVIDIA_API_KEY"] = nvidai_api_key
+    from agno.models.message import Message
 
-    # Create Nvidia instance with conditional base_url
-    kwargs = {"id": config_obj.get("id")}
-    if "base_url" in config_obj and config_obj.get("base_url") is not None:
-        kwargs["base_url"] = config_obj.get("base_url")
+    class AgnoThinkingInjector(BaseThinkingInjector):
 
-    client = Nvidia(**kwargs)  # type: ignore[arg-type]
+        from agno.models.message import Message
 
-    if isinstance(client, RetryMixin):
+        @override
+        def inject(self, messages: list[Message], *args, **kwargs) -> FunctionArgumentWrapper:
+            new_messages = [Message(role="system", content=self.system_prompt)] + messages
+            return FunctionArgumentWrapper(new_messages, *args, **kwargs)
 
+    if isinstance(llm_config, RetryMixin):
         client = patch_with_retry(client,
                                   retries=llm_config.num_retries,
                                   retry_codes=llm_config.retry_on_status_codes,
                                   retry_on_messages=llm_config.retry_on_errors)
 
-    yield client
+    if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
+        client = patch_with_thinking(
+            client,
+            AgnoThinkingInjector(system_prompt=llm_config.thinking_system_prompt,
+                                 function_names=[
+                                     "invoke_stream",
+                                     "invoke",
+                                     "ainvoke",
+                                     "ainvoke_stream",
+                                 ]))
 
+    return client
 
-@register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.AGNO)
-async def openai_agno(llm_config: OpenAIModelConfig, builder: Builder):
 
-    from agno.models.openai import OpenAIChat
+@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.AGNO)
+async def nim_agno(llm_config: NIMModelConfig, _builder: Builder):
+
+    from agno.models.nvidia import Nvidia
+
+    config_obj = {
+        **llm_config.model_dump(
+            exclude={"type", "model_name", "thinking"},
+            by_alias=True,
+            exclude_none=True,
+        ),
+    }
 
-    # Use model_dump to get the proper field values with correct types
-    kwargs = llm_config.model_dump(exclude={"type"}, by_alias=True)
+    client = Nvidia(**config_obj, id=llm_config.model_name)
 
-    # AGNO uses 'id' instead of 'model' for the model name
-    if "model" in kwargs:
-        kwargs["id"] = kwargs.pop("model")
+    yield _patch_llm_based_on_config(client, llm_config)
 
-    client = OpenAIChat(**kwargs)
 
-    if isinstance(llm_config, RetryMixin):
-        client = patch_with_retry(client,
-                                  retries=llm_config.num_retries,
-                                  retry_codes=llm_config.retry_on_status_codes,
-                                  retry_on_messages=llm_config.retry_on_errors)
+@register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.AGNO)
+async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder):
+
+    from agno.models.openai import OpenAIChat
+
+    config_obj = {
+        **llm_config.model_dump(
+            exclude={"type", "model_name", "thinking"},
+            by_alias=True,
+            exclude_none=True,
+        ),
+    }
+
+    client = OpenAIChat(**config_obj, id=llm_config.model_name)
 
-    yield client
+    yield _patch_llm_based_on_config(client, llm_config)

nat/plugins/agno/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 

nat/plugins/agno/tool_wrapper.py

@@ -17,11 +17,9 @@ import asyncio
 import json
 import logging
 import textwrap
-import traceback
+from collections.abc import Awaitable
+from collections.abc import Callable
 from typing import Any
-from typing import Awaitable
-from typing import Callable
-from typing import List
 
 from agno.tools import tool
 
@@ -134,7 +132,7 @@ async def process_result(result: Any, name: str) -> str:
 
 def execute_agno_tool(name: str,
                       coroutine_fn: Callable[..., Awaitable[Any]],
-                      required_fields: List[str],
+                      required_fields: list[str],
                       loop: asyncio.AbstractEventLoop,
                       **kwargs: Any) -> Any:
     """
@@ -146,7 +144,7 @@ def execute_agno_tool(name: str,
         The name of the tool
     coroutine_fn : Callable
         The async function to invoke
-    required_fields : List[str]
+    required_fields : list[str]
         List of required fields for validation
     loop : asyncio.AbstractEventLoop
         The event loop to use for async execution
@@ -157,7 +155,6 @@ def execute_agno_tool(name: str,
     -------
     The result of the function execution as a string
     """
-    global _tool_call_counters, _tool_initialization_done
 
     try:
         logger.debug(f"Running {name} with kwargs: {kwargs}")
@@ -288,9 +285,7 @@ def execute_agno_tool(name: str,
         return process_future.result(timeout=30)  # 30-second timeout for processing
 
     except Exception as e:
-        logger.exception(f"Error executing Agno tool {name}: {e}")
-        error_traceback = traceback.format_exc()
-        logger.error(f"Exception traceback: {error_traceback}")
+        logger.error("Error executing Agno tool %s: %s", name, e)
         raise
 
 
@@ -335,7 +330,7 @@ def agno_tool_wrapper(name: str, fn: Function, builder: Builder):
     if description:
         description = textwrap.dedent(description).strip()
 
-    # Input schema handling from LangChain-style
+    # Input schema handling from LangChain/LangGraph-style
     required_fields = []
     if fn.input_schema is not None:
         try:

nat/plugins/agno/tools/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 

{nvidia_nat_agno-1.3.dev0.dist-info → nvidia_nat_agno-1.3.0rc1.dist-info}/METADATA

@@ -1,14 +1,16 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-agno
-Version: 1.3.dev0
+Version: 1.3.0rc1
 Summary: Subpackage for Agno integration in NeMo Agent toolkit
 Keywords: ai,rag,agents
 Classifier: Programming Language :: Python
-Requires-Python: <3.13,>=3.11
+Classifier: Programming Language :: Python :: 3.11
+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-dev
+Requires-Dist: nvidia-nat[openai]==v1.3.0-rc1
 Requires-Dist: agno~=1.2.3
-Requires-Dist: openai~=1.66
 Requires-Dist: google-search-results~=2.4.2
 
 <!--

nvidia_nat_agno-1.3.0rc1.dist-info/RECORD

@@ -0,0 +1,13 @@
+nat/meta/pypi.md,sha256=tZD7hiOSYWgiAdddD1eIJ8T5ipZwEIjnd8ilgmasdmw,1198
+nat/plugins/agno/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nat/plugins/agno/llm.py,sha256=op02DUgf_DwYFxOwmUmd7NhFIf6TTXj313sFV6GOZZM,3828
+nat/plugins/agno/register.py,sha256=q-es1KVb_PTHOFszltym7e7Pj2jmyieih_Ve-cguHI8,788
+nat/plugins/agno/tool_wrapper.py,sha256=Iq8v0uO6HfSVQvcRwZOk4TMRpRyWUMM4s-XAtZx_Hj8,15729
+nat/plugins/agno/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nat/plugins/agno/tools/register.py,sha256=mpNgD1r51EoYlQiEAqPX15wCtx2alDcMIQiR8F_ub-M,743
+nat/plugins/agno/tools/serp_api_tool.py,sha256=AJQH6-1iEUUrk_nzfZ3zZqutEKhJ_LMOUJi_iol65Sc,4442
+nvidia_nat_agno-1.3.0rc1.dist-info/METADATA,sha256=8EKugYRK44pr-1SWNsOJGiobrJwL86I_iqcVD02yBAY,1731
+nvidia_nat_agno-1.3.0rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_agno-1.3.0rc1.dist-info/entry_points.txt,sha256=qRhuHKj2WmdJkLpbVXpYkdtc2cZdG4LPlBsABG2ImVI,103
+nvidia_nat_agno-1.3.0rc1.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_agno-1.3.0rc1.dist-info/RECORD,,

nvidia_nat_agno-1.3.dev0.dist-info/RECORD

@@ -1,13 +0,0 @@
-nat/meta/pypi.md,sha256=tZD7hiOSYWgiAdddD1eIJ8T5ipZwEIjnd8ilgmasdmw,1198
-nat/plugins/agno/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nat/plugins/agno/llm.py,sha256=hvPu0-gkRasnWZY5kwHOiA6rZ1ZIGlnaH3nPPfqrFJ8,3434
-nat/plugins/agno/register.py,sha256=6vC1TjMxo3igqTnEtVFgLEf_jgLYkBfBZxjwqxGng6w,820
-nat/plugins/agno/tool_wrapper.py,sha256=uClYG-LvRj1OBTMwzMuE40cibLhqVRXSCiqniVQND2Y,15914
-nat/plugins/agno/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nat/plugins/agno/tools/register.py,sha256=OCmzR03CHmQHm34ZEascM1dRVh-ALMs2mafDcqLDz6s,775
-nat/plugins/agno/tools/serp_api_tool.py,sha256=AJQH6-1iEUUrk_nzfZ3zZqutEKhJ_LMOUJi_iol65Sc,4442
-nvidia_nat_agno-1.3.dev0.dist-info/METADATA,sha256=uLs9BmqF2hQBiwk05Mqi4lKFDwQSS5EkgQ8Q4zwcKfE,1596
-nvidia_nat_agno-1.3.dev0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_agno-1.3.dev0.dist-info/entry_points.txt,sha256=qRhuHKj2WmdJkLpbVXpYkdtc2cZdG4LPlBsABG2ImVI,103
-nvidia_nat_agno-1.3.dev0.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_agno-1.3.dev0.dist-info/RECORD,,