nvidia-nat-llama-index 1.1.0a20251020__py3-none-any.whl

nat/meta/pypi.md

@@ -0,0 +1,23 @@
+<!--
+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.
+-->
+
+![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/banner.png "NeMo Agent toolkit banner image")
+
+# NVIDIA NeMo Agent Toolkit Subpackage
+This is a subpackage for Llama-Index integration in NeMo Agent toolkit.
+
+For more information about the NVIDIA NeMo Agent toolkit, please visit the [NeMo Agent toolkit GitHub Repo](https://github.com/NVIDIA/NeMo-Agent-Toolkit).

nat/plugins/llama_index/embedder.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_embedder_client
+from nat.data_models.retry_mixin import RetryMixin
+from nat.embedder.azure_openai_embedder import AzureOpenAIEmbedderModelConfig
+from nat.embedder.nim_embedder import NIMEmbedderModelConfig
+from nat.embedder.openai_embedder import OpenAIEmbedderModelConfig
+from nat.utils.exception_handlers.automatic_retries import patch_with_retry
+
+
+@register_embedder_client(config_type=AzureOpenAIEmbedderModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def azure_openai_llama_index(embedder_config: AzureOpenAIEmbedderModelConfig, _builder: Builder):
+
+    from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
+
+    client = AzureOpenAIEmbedding(**embedder_config.model_dump(exclude={"type"}, by_alias=True, exclude_none=True))
+
+    if isinstance(embedder_config, RetryMixin):
+        client = patch_with_retry(client,
+                                  retries=embedder_config.num_retries,
+                                  retry_codes=embedder_config.retry_on_status_codes,
+                                  retry_on_messages=embedder_config.retry_on_errors)
+
+    yield client
+
+
+@register_embedder_client(config_type=NIMEmbedderModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def nim_llama_index(embedder_config: NIMEmbedderModelConfig, _builder: Builder):
+
+    from llama_index.embeddings.nvidia import NVIDIAEmbedding  # pylint: disable=no-name-in-module
+
+    client = NVIDIAEmbedding(
+        **embedder_config.model_dump(exclude={"type", "model_name"}, by_alias=True, exclude_none=True),
+        model=embedder_config.model_name,
+    )
+
+    if isinstance(embedder_config, RetryMixin):
+        client = patch_with_retry(client,
+                                  retries=embedder_config.num_retries,
+                                  retry_codes=embedder_config.retry_on_status_codes,
+                                  retry_on_messages=embedder_config.retry_on_errors)
+
+    yield client
+
+
+@register_embedder_client(config_type=OpenAIEmbedderModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def openai_llama_index(embedder_config: OpenAIEmbedderModelConfig, _builder: Builder):
+
+    from llama_index.embeddings.openai import OpenAIEmbedding
+
+    client = OpenAIEmbedding(**embedder_config.model_dump(exclude={"type"}, by_alias=True, exclude_none=True))
+
+    if isinstance(embedder_config, RetryMixin):
+        client = patch_with_retry(client,
+                                  retries=embedder_config.num_retries,
+                                  retry_codes=embedder_config.retry_on_status_codes,
+                                  retry_on_messages=embedder_config.retry_on_errors)
+
+    yield client

nat/plugins/llama_index/llm.py

@@ -0,0 +1,128 @@
+# 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.
+
+from collections.abc import Sequence
+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.aws_bedrock_llm import AWSBedrockModelConfig
+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
+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")
+
+
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+
+    from llama_index.core.base.llms.types import ChatMessage
+
+    class LlamaIndexThinkingInjector(BaseThinkingInjector):
+
+        @override
+        def inject(self, messages: Sequence[ChatMessage], *args, **kwargs) -> FunctionArgumentWrapper:
+            for i, message in enumerate(messages):
+                if message.role == "system":
+                    if self.system_prompt not in str(message.content):
+                        messages = list(messages)
+                        messages[i] = ChatMessage(role="system", content=f"{message.content}\n{self.system_prompt}")
+                    break
+            else:
+                messages = list(messages)
+                messages.insert(0, ChatMessage(role="system", content=self.system_prompt))
+            return FunctionArgumentWrapper(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)
+
+    if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
+        client = patch_with_thinking(
+            client,
+            LlamaIndexThinkingInjector(
+                system_prompt=llm_config.thinking_system_prompt,
+                function_names=[
+                    "chat",
+                    "stream_chat",
+                    "achat",
+                    "astream_chat",
+                ],
+            ))
+
+    return client
+
+
+@register_llm_client(config_type=AWSBedrockModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def aws_bedrock_llama_index(llm_config: AWSBedrockModelConfig, _builder: Builder):
+
+    from llama_index.llms.bedrock import Bedrock
+
+    # LlamaIndex uses context_size instead of max_tokens
+    llm = Bedrock(**llm_config.model_dump(exclude={"type", "top_p", "thinking"}, by_alias=True))
+
+    yield _patch_llm_based_on_config(llm, llm_config)
+
+
+@register_llm_client(config_type=AzureOpenAIModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def azure_openai_llama_index(llm_config: AzureOpenAIModelConfig, _builder: Builder):
+
+    from llama_index.llms.azure_openai import AzureOpenAI
+
+    llm = AzureOpenAI(**llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True))
+
+    yield _patch_llm_based_on_config(llm, llm_config)
+
+
+@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def nim_llama_index(llm_config: NIMModelConfig, _builder: Builder):
+
+    from llama_index.llms.nvidia import NVIDIA
+
+    llm = NVIDIA(**llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True))
+
+    yield _patch_llm_based_on_config(llm, llm_config)
+
+
+@register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def openai_llama_index(llm_config: OpenAIModelConfig, _builder: Builder):
+
+    from llama_index.llms.openai import OpenAI
+
+    llm = OpenAI(**llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True))
+
+    yield _patch_llm_based_on_config(llm, llm_config)
+
+
+@register_llm_client(config_type=LiteLlmModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def litellm_llama_index(llm_config: LiteLlmModelConfig, _builder: Builder):
+
+    from llama_index.llms.litellm import LiteLLM
+
+    llm = LiteLLM(**llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True))
+
+    yield _patch_llm_based_on_config(llm, llm_config)

nat/plugins/llama_index/register.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.
+
+# flake8: noqa
+# isort:skip_file
+
+# Import any providers which need to be automatically registered here
+
+from . import embedder
+from . import llm
+from . import tool_wrapper

nat/plugins/llama_index/tool_wrapper.py

@@ -0,0 +1,32 @@
+# 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.
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function import Function
+from nat.cli.register_workflow import register_tool_wrapper
+
+
+@register_tool_wrapper(wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+def langchain_tool_wrapper(name: str, fn: Function, builder: Builder):
+
+    from llama_index.core.tools import FunctionTool
+
+    assert fn.input_schema is not None, "Tool must have input schema"
+
+    return FunctionTool.from_defaults(async_fn=fn.acall_invoke,
+                                      name=name,
+                                      description=fn.description,
+                                      fn_schema=fn.input_schema)

nvidia_nat_llama_index-1.1.0a20251020.dist-info/METADATA

@@ -0,0 +1,55 @@
+Metadata-Version: 2.4
+Name: nvidia-nat-llama-index
+Version: 1.1.0a20251020
+Summary: Subpackage for Llama-Index integration in NeMo Agent toolkit
+Author: NVIDIA Corporation
+Maintainer: NVIDIA Corporation
+License: Apache-2.0
+Project-URL: documentation, https://docs.nvidia.com/nemo/agent-toolkit/latest/
+Project-URL: source, https://github.com/NVIDIA/NeMo-Agent-Toolkit
+Keywords: ai,rag,agents
+Classifier: Programming Language :: Python
+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
+License-File: LICENSE-3rd-party.txt
+License-File: LICENSE.md
+Requires-Dist: nvidia-nat==v1.1.0a20251020
+Requires-Dist: llama-index-core~=0.12.21
+Requires-Dist: llama-index-embeddings-azure-openai~=0.3.9
+Requires-Dist: llama-index-embeddings-nvidia~=0.3.1
+Requires-Dist: llama-index-embeddings-openai~=0.3.1
+Requires-Dist: llama-index-llms-azure-openai~=0.3.2
+Requires-Dist: llama-index-llms-bedrock~=0.3.8
+Requires-Dist: llama-index-llms-litellm~=0.5.1
+Requires-Dist: llama-index-llms-nvidia~=0.3.1
+Requires-Dist: llama-index-llms-openai~=0.3.42
+Requires-Dist: llama-index-readers-file~=0.4.4
+Requires-Dist: llama-index~=0.12.21
+Dynamic: license-file
+
+<!--
+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.
+-->
+
+![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/banner.png "NeMo Agent toolkit banner image")
+
+# NVIDIA NeMo Agent Toolkit Subpackage
+This is a subpackage for Llama-Index integration in NeMo Agent toolkit.
+
+For more information about the NVIDIA NeMo Agent toolkit, please visit the [NeMo Agent toolkit GitHub Repo](https://github.com/NVIDIA/NeMo-Agent-Toolkit).

nvidia_nat_llama_index-1.1.0a20251020.dist-info/RECORD

@@ -0,0 +1,13 @@
+nat/meta/pypi.md,sha256=s9C3pgWB0HLIXTx5QPryNOWN0O2fIRIap0p9_zCHlTs,1112
+nat/plugins/llama_index/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nat/plugins/llama_index/embedder.py,sha256=4c4uAa19UArqA_Npe5mwEnUhCaguTbyWcAntBzcPEeQ,3493
+nat/plugins/llama_index/llm.py,sha256=7fneFsjXjZJ-ReLAy1aUslkGhj6_fG_gFbyy4ME-QH4,5350
+nat/plugins/llama_index/register.py,sha256=1x_b8u6cuQwh4Iz_7TcIFWXvLIL9IIKUPE-zR9d6ug8,859
+nat/plugins/llama_index/tool_wrapper.py,sha256=VFKMIIeLdWqHwW2Ax11E2w-_9w3ow6Iuhra1Hk78RYM,1387
+nvidia_nat_llama_index-1.1.0a20251020.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
+nvidia_nat_llama_index-1.1.0a20251020.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+nvidia_nat_llama_index-1.1.0a20251020.dist-info/METADATA,sha256=Nd0eKWPPvyiacUcqYPWndvRt9uM8-TohPBvhmGwjujY,2422
+nvidia_nat_llama_index-1.1.0a20251020.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_llama_index-1.1.0a20251020.dist-info/entry_points.txt,sha256=2LqRRju5448P2v8B3y6TSPnk-nOd5T3AmV5JibCnoQc,68
+nvidia_nat_llama_index-1.1.0a20251020.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_llama_index-1.1.0a20251020.dist-info/RECORD,,

nvidia_nat_llama_index-1.1.0a20251020.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

nvidia_nat_llama_index-1.1.0a20251020.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[nat.components]
+nat_llama_index = nat.plugins.llama_index.register