nvidia-nat-llama-index 1.3.0a20250826__py3-none-any.whl → 1.3.0a20250828__py3-none-any.whl

nat/plugins/llama_index/llm.py

@@ -13,33 +13,70 @@
 # 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.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=AWSBedrockModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
-async def aws_bedrock_llama_index(llm_config: AWSBedrockModelConfig, _builder: Builder):
 
-    from llama_index.llms.bedrock import Bedrock
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
 
-    kwargs = llm_config.model_dump(exclude={"type", "max_tokens"}, by_alias=True)
+    from llama_index.core.base.llms.types import ChatMessage
 
-    llm = Bedrock(**kwargs)
+    class LlamaIndexThinkingInjector(BaseThinkingInjector):
+
+        @override
+        def inject(self, messages: Sequence[ChatMessage], *args, **kwargs) -> FunctionArgumentWrapper:
+            new_messages = [ChatMessage(role="system", content=self.system_prompt)] + list(messages)
+            return FunctionArgumentWrapper(new_messages, *args, **kwargs)
+
+    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",
+                ],
+            ))
 
     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)
+        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)
+
+    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
 
-    yield llm
+    # LlamaIndex uses context_size instead of max_tokens
+    llm = Bedrock(**llm_config.model_dump(exclude={"type", "top_p"}, by_alias=True), )
+
+    yield _patch_llm_based_on_config(llm, llm_config)
 
 
 @register_llm_client(config_type=AzureOpenAIModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
@@ -47,17 +84,9 @@ async def azure_openai_llama_index(llm_config: AzureOpenAIModelConfig, _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,
-                               retries=llm_config.num_retries,
-                               retry_codes=llm_config.retry_on_status_codes,
-                               retry_on_messages=llm_config.retry_on_errors)
+    llm = AzureOpenAI(**llm_config.model_dump(exclude={"type"}, by_alias=True))
 
-    yield llm
+    yield _patch_llm_based_on_config(llm, llm_config)
 
 
 @register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
@@ -65,20 +94,9 @@ async def nim_llama_index(llm_config: NIMModelConfig, _builder: Builder):
 
     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 = NVIDIA(**kwargs)
+    llm = NVIDIA(**llm_config.model_dump(exclude={"type"}, by_alias=True, exclude_none=True))
 
-    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
+    yield _patch_llm_based_on_config(llm, llm_config)
 
 
 @register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
@@ -86,17 +104,6 @@ async def openai_llama_index(llm_config: OpenAIModelConfig, _builder: Builder):
 
     from llama_index.llms.openai import OpenAI
 
-    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)
-
-    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)
+    llm = OpenAI(**llm_config.model_dump(exclude={"type"}, by_alias=True, exclude_none=True))
 
-    yield llm
+    yield _patch_llm_based_on_config(llm, llm_config)

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

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-llama-index
-Version: 1.3.0a20250826
+Version: 1.3.0a20250828
 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.0a20250826
+Requires-Dist: nvidia-nat==v1.3.0a20250828
 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

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

@@ -1,11 +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=6US_DR4dWwHuhRnf2KlN8TVIV7wk7NHTDFyBwgMyEQc,3093
-nat/plugins/llama_index/llm.py,sha256=PtXibAiZWCv1pjEeP5A77Tyi-rVUGCHc81jEoIdy3p8,3949
+nat/plugins/llama_index/llm.py,sha256=_aQ1hgM-YR7GleotwlSe6Dw4Of0wdwKRLCYZMN2VcIU,4454
 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.3.0a20250826.dist-info/METADATA,sha256=FmUGKqn6BuxmgR9v0Z-xaIOBe-U2iIENV-CnUG9ox0o,1916
-nvidia_nat_llama_index-1.3.0a20250826.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_llama_index-1.3.0a20250826.dist-info/entry_points.txt,sha256=2LqRRju5448P2v8B3y6TSPnk-nOd5T3AmV5JibCnoQc,68
-nvidia_nat_llama_index-1.3.0a20250826.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_llama_index-1.3.0a20250826.dist-info/RECORD,,
+nvidia_nat_llama_index-1.3.0a20250828.dist-info/METADATA,sha256=kOatcvdy8DXcTz8hVEsiU84unXu3vcgFZ9By_bY6G_M,1916
+nvidia_nat_llama_index-1.3.0a20250828.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_llama_index-1.3.0a20250828.dist-info/entry_points.txt,sha256=2LqRRju5448P2v8B3y6TSPnk-nOd5T3AmV5JibCnoQc,68
+nvidia_nat_llama_index-1.3.0a20250828.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_llama_index-1.3.0a20250828.dist-info/RECORD,,