nvidia-nat-llama-index 1.3.dev0__py3-none-any.whl → 1.3.0.dev2__py3-none-any.whl

nat/plugins/llama_index/embedder.py

@@ -12,18 +12,35 @@
 # 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-argument
 
 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))
+
+    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_llamaindex(embedder_config: NIMEmbedderModelConfig, builder: Builder):
+async def nim_llamaindex(embedder_config: NIMEmbedderModelConfig, _builder: Builder):
 
     from llama_index.embeddings.nvidia import NVIDIAEmbedding  # pylint: disable=no-name-in-module
 
@@ -33,7 +50,15 @@ async def nim_llamaindex(embedder_config: NIMEmbedderModelConfig, builder: Build
             embedder_config.model_name,
     }
 
-    client = NVIDIAEmbedding(**config_obj)
+    yield NVIDIAEmbedding(**config_obj)
+
+
+@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))
 
     if isinstance(embedder_config, RetryMixin):
         client = patch_with_retry(client,

nat/plugins/llama_index/llm.py

@@ -18,22 +18,38 @@ from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.register_workflow import register_llm_client
 from nat.data_models.retry_mixin import RetryMixin
 from nat.llm.aws_bedrock_llm import AWSBedrockModelConfig
+from nat.llm.azure_openai_llm import AzureOpenAIModelConfig
 from nat.llm.nim_llm import NIMModelConfig
 from nat.llm.openai_llm import OpenAIModelConfig
 from nat.utils.exception_handlers.automatic_retries import patch_with_retry
 
 
-@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
-async def nim_llama_index(llm_config: NIMModelConfig, builder: Builder):
+@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.nvidia import NVIDIA
+    from llama_index.llms.bedrock import Bedrock
 
-    kwargs = llm_config.model_dump(exclude={"type"}, by_alias=True)
+    kwargs = llm_config.model_dump(exclude={"type", "max_tokens"}, by_alias=True)
 
-    if ("base_url" in kwargs and kwargs["base_url"] is None):
-        del kwargs["base_url"]
+    llm = Bedrock(**kwargs)
 
-    llm = NVIDIA(**kwargs)
+    if isinstance(llm_config, RetryMixin):
+        llm = patch_with_retry(llm,
+                               retries=llm_config.num_retries,
+                               retry_codes=llm_config.retry_on_status_codes,
+                               retry_on_messages=llm_config.retry_on_errors)
+
+    yield llm
+
+
+@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
+
+    kwargs = llm_config.model_dump(exclude={"type"}, by_alias=True)
+
+    llm = AzureOpenAI(**kwargs)
 
     if isinstance(llm_config, RetryMixin):
         llm = patch_with_retry(llm,
@@ -44,17 +60,17 @@ async def nim_llama_index(llm_config: NIMModelConfig, builder: Builder):
     yield llm
 
 
-@register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
-async def openai_llama_index(llm_config: OpenAIModelConfig, builder: Builder):
+@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.openai import OpenAI
+    from llama_index.llms.nvidia import NVIDIA
 
     kwargs = llm_config.model_dump(exclude={"type"}, by_alias=True)
 
     if ("base_url" in kwargs and kwargs["base_url"] is None):
         del kwargs["base_url"]
 
-    llm = OpenAI(**kwargs)
+    llm = NVIDIA(**kwargs)
 
     if isinstance(llm_config, RetryMixin):
         llm = patch_with_retry(llm,
@@ -65,14 +81,17 @@ async def openai_llama_index(llm_config: OpenAIModelConfig, builder: Builder):
     yield llm
 
 
-@register_llm_client(config_type=AWSBedrockModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
-async def aws_bedrock_llama_index(llm_config: AWSBedrockModelConfig, builder: Builder):
+@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.bedrock import Bedrock
+    from llama_index.llms.openai import OpenAI
 
-    kwargs = llm_config.model_dump(exclude={"type", "max_tokens"}, by_alias=True)
+    kwargs = llm_config.model_dump(exclude={"type"}, by_alias=True)
 
-    llm = Bedrock(**kwargs)
+    if ("base_url" in kwargs and kwargs["base_url"] is None):
+        del kwargs["base_url"]
+
+    llm = OpenAI(**kwargs)
 
     if isinstance(llm_config, RetryMixin):
         llm = patch_with_retry(llm,

nat/plugins/llama_index/register.py

@@ -19,6 +19,6 @@
 
 # Import any providers which need to be automatically registered here
 
+from . import embedder
 from . import llm
 from . import tool_wrapper
-from . import embedder

{nvidia_nat_llama_index-1.3.dev0.dist-info → nvidia_nat_llama_index-1.3.0.dev2.dist-info}/METADATA

@@ -1,18 +1,22 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-llama-index
-Version: 1.3.dev0
+Version: 1.3.0.dev2
 Summary: Subpackage for Llama-Index 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-dev
-Requires-Dist: llama-index-core==0.12.21
-Requires-Dist: llama-index-embeddings-nvidia==0.3.1
-Requires-Dist: llama-index-llms-bedrock==0.3.8
-Requires-Dist: llama-index-llms-nvidia==0.3.1
-Requires-Dist: llama-index-readers-file==0.4.4
-Requires-Dist: llama-index==0.12.21
+Requires-Dist: nvidia-nat==v1.3.0-dev2
+Requires-Dist: llama-index-core~=0.12.21
+Requires-Dist: llama-index-embeddings-azure-openai~=0.3.1
+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.1
+Requires-Dist: llama-index-llms-bedrock~=0.3.8
+Requires-Dist: llama-index-llms-nvidia~=0.3.1
+Requires-Dist: llama-index-llms-openai~=0.3.38
+Requires-Dist: llama-index-readers-file~=0.4.4
+Requires-Dist: llama-index~=0.12.21
 
 <!--
 SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

nvidia_nat_llama_index-1.3.0.dev2.dist-info/RECORD

@@ -0,0 +1,11 @@
+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=ROGlJJxwIZXV3Sg93Z9RIelOOhxuJcrw0B1id43tn_U,3130
+nat/plugins/llama_index/llm.py,sha256=PtXibAiZWCv1pjEeP5A77Tyi-rVUGCHc81jEoIdy3p8,3949
+nat/plugins/llama_index/register.py,sha256=ICmUCqrNvWTDFe6Zjm6B6KUkgqsyqfFPTHMGU74KsA4,891
+nat/plugins/llama_index/tool_wrapper.py,sha256=VFKMIIeLdWqHwW2Ax11E2w-_9w3ow6Iuhra1Hk78RYM,1387
+nvidia_nat_llama_index-1.3.0.dev2.dist-info/METADATA,sha256=uKukv46PucXHXNFo51uQa_FkYeR6JHtk91E5W8cIzxY,1908
+nvidia_nat_llama_index-1.3.0.dev2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_llama_index-1.3.0.dev2.dist-info/entry_points.txt,sha256=2LqRRju5448P2v8B3y6TSPnk-nOd5T3AmV5JibCnoQc,68
+nvidia_nat_llama_index-1.3.0.dev2.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_llama_index-1.3.0.dev2.dist-info/RECORD,,

nvidia_nat_llama_index-1.3.dev0.dist-info/RECORD

@@ -1,11 +0,0 @@
-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=L7Wx9QJH2Sh0y6kmgisjtTZfDC9YcuuusKl0Tn_7JgA,1921
-nat/plugins/llama_index/llm.py,sha256=VyuIqks2G9Zcs9a0X7jJyASzSUdtL0Pns0v4_rFEsJE,3220
-nat/plugins/llama_index/register.py,sha256=DelI6uth8dO8h_7czPW4Akfr_ajXDBwhEa1dJzRf5u4,891
-nat/plugins/llama_index/tool_wrapper.py,sha256=VFKMIIeLdWqHwW2Ax11E2w-_9w3ow6Iuhra1Hk78RYM,1387
-nvidia_nat_llama_index-1.3.dev0.dist-info/METADATA,sha256=EP1nCqjUnjLBuvlh_JKpFIG_rebxwmKzo0MjDejM0MY,1694
-nvidia_nat_llama_index-1.3.dev0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_llama_index-1.3.dev0.dist-info/entry_points.txt,sha256=2LqRRju5448P2v8B3y6TSPnk-nOd5T3AmV5JibCnoQc,68
-nvidia_nat_llama_index-1.3.dev0.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_llama_index-1.3.dev0.dist-info/RECORD,,