nvidia-nat-agno 1.3.0a20250826__py3-none-any.whl → 1.3.0a20250828__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
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
 
-    config_obj = {
-        **llm_config.model_dump(exclude={"type", "model_name"}, by_alias=True),
-        "id": f"{llm_config.model_name}",
-    }
+    from agno.models.message import Message
 
-    # 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):
+    class AgnoThinkingInjector(BaseThinkingInjector):
 
-        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")
+        from agno.models.message import Message
 
-            if (nvidai_api_key is not None):
-                # Transfer the key to the correct environment variable
-                os.environ["NVIDIA_API_KEY"] = nvidai_api_key
+        @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)
 
-    # 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")
-
-    client = Nvidia(**kwargs)  # type: ignore[arg-type]
-
-    if isinstance(client, RetryMixin):
+    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",
+                                 ]))
 
+    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
+    return client
 
 
-@register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.AGNO)
-async def openai_agno(llm_config: OpenAIModelConfig, builder: Builder):
+@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.AGNO)
+async def nim_agno(llm_config: NIMModelConfig, _builder: Builder):
 
-    from agno.models.openai import OpenAIChat
+    from agno.models.nvidia import Nvidia
+
+    config_obj = {
+        **llm_config.model_dump(
+            exclude={"type", "model_name"},
+            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"},
+            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)

{nvidia_nat_agno-1.3.0a20250826.dist-info → nvidia_nat_agno-1.3.0a20250828.dist-info}/METADATA

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-agno
-Version: 1.3.0a20250826
+Version: 1.3.0a20250828
 Summary: Subpackage for Agno integration in NeMo Agent toolkit
 Keywords: ai,rag,agents
 Classifier: Programming Language :: Python
 Requires-Python: <3.13,>=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: nvidia-nat==v1.3.0a20250826
+Requires-Dist: nvidia-nat==v1.3.0a20250828
 Requires-Dist: agno~=1.2.3
 Requires-Dist: openai~=1.66
 Requires-Dist: google-search-results~=2.4.2

{nvidia_nat_agno-1.3.0a20250826.dist-info → nvidia_nat_agno-1.3.0a20250828.dist-info}/RECORD

@@ -1,13 +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=hvPu0-gkRasnWZY5kwHOiA6rZ1ZIGlnaH3nPPfqrFJ8,3434
+nat/plugins/agno/llm.py,sha256=QWJAQopnsb51efVkTHYsyolvkBB09dFSLpwjn847qTY,3804
 nat/plugins/agno/register.py,sha256=q-es1KVb_PTHOFszltym7e7Pj2jmyieih_Ve-cguHI8,788
 nat/plugins/agno/tool_wrapper.py,sha256=VEslkY_jerrgZDyVnW5WSjC3xl2gt8-afSMODZ1HMzU,15856
 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.0a20250826.dist-info/METADATA,sha256=PQu4UAZSW7251-3tlC_n9PN_cP6Lt09aS4SDmzMpi34,1609
-nvidia_nat_agno-1.3.0a20250826.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_agno-1.3.0a20250826.dist-info/entry_points.txt,sha256=qRhuHKj2WmdJkLpbVXpYkdtc2cZdG4LPlBsABG2ImVI,103
-nvidia_nat_agno-1.3.0a20250826.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_agno-1.3.0a20250826.dist-info/RECORD,,
+nvidia_nat_agno-1.3.0a20250828.dist-info/METADATA,sha256=9TeTsXNymJBQapUI9eBcFtT9nVsm_X0GvXc7tpRbcLI,1609
+nvidia_nat_agno-1.3.0a20250828.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_agno-1.3.0a20250828.dist-info/entry_points.txt,sha256=qRhuHKj2WmdJkLpbVXpYkdtc2cZdG4LPlBsABG2ImVI,103
+nvidia_nat_agno-1.3.0a20250828.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_agno-1.3.0a20250828.dist-info/RECORD,,