nvidia-nat-langchain 1.3.0a20250826__py3-none-any.whl → 1.3.0a20250828__py3-none-any.whl

nat/plugins/langchain/llm.py

@@ -13,15 +13,96 @@
 # 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")
+
+
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+
+    from langchain_core.language_models import LanguageModelInput
+    from langchain_core.messages import BaseMessage
+    from langchain_core.messages import HumanMessage
+    from langchain_core.messages import SystemMessage
+    from langchain_core.prompt_values import PromptValue
+
+    class LangchainThinkingInjector(BaseThinkingInjector):
+
+        @override
+        def inject(self, messages: LanguageModelInput, *args, **kwargs) -> FunctionArgumentWrapper:
+            """
+            Inject a system prompt into the messages.
+
+            The messages are the first (non-object) argument to the function.
+            The rest of the arguments are passed through unchanged.
+
+            Args:
+                messages: The messages to inject the system prompt into.
+                *args: The rest of the arguments to the function.
+                **kwargs: The rest of the keyword arguments to the function.
+
+            Returns:
+                FunctionArgumentWrapper: An object that contains the transformed args and kwargs.
+
+            Raises:
+                ValueError: If the messages are not a valid type for LanguageModelInput.
+            """
+            system_message = SystemMessage(content=self.system_prompt)
+            if isinstance(messages, BaseMessage):
+                new_messages = [system_message, messages]
+                return FunctionArgumentWrapper(new_messages, *args, **kwargs)
+            elif isinstance(messages, PromptValue):
+                new_messages = [system_message, *messages.to_messages()]
+                return FunctionArgumentWrapper(new_messages, *args, **kwargs)
+            elif isinstance(messages, str):
+                new_messages = [system_message, HumanMessage(content=messages)]
+                return FunctionArgumentWrapper(new_messages, *args, **kwargs)
+            elif isinstance(messages, Sequence):
+                if all(isinstance(m, BaseMessage) for m in messages):
+                    new_messages = [system_message, *list(messages)]
+                    return FunctionArgumentWrapper(new_messages, *args, **kwargs)
+                raise ValueError(
+                    "Unsupported sequence element types for LanguageModelInput; expected Sequence[BaseMessage].")
+            else:
+                return FunctionArgumentWrapper(messages, *args, **kwargs)
+
+    if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
+        client = patch_with_thinking(
+            client,
+            LangchainThinkingInjector(
+                system_prompt=llm_config.thinking_system_prompt,
+                function_names=[
+                    "invoke",
+                    "ainvoke",
+                    "stream",
+                    "astream",
+                ],
+            ))
+
+    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)
+
+    return client
 
 
 @register_llm_client(config_type=AWSBedrockModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
@@ -31,13 +112,7 @@ async def aws_bedrock_langchain(llm_config: AWSBedrockModelConfig, _builder: Bui
 
     client = ChatBedrockConverse(**llm_config.model_dump(exclude={"type", "context_size"}, by_alias=True))
 
-    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
+    yield _patch_llm_based_on_config(client, llm_config)
 
 
 @register_llm_client(config_type=AzureOpenAIModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
@@ -47,13 +122,7 @@ async def azure_openai_langchain(llm_config: AzureOpenAIModelConfig, _builder: B
 
     client = AzureChatOpenAI(**llm_config.model_dump(exclude={"type"}, by_alias=True))
 
-    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
+    yield _patch_llm_based_on_config(client, llm_config)
 
 
 @register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
@@ -61,15 +130,13 @@ async def nim_langchain(llm_config: NIMModelConfig, _builder: Builder):
 
     from langchain_nvidia_ai_endpoints import ChatNVIDIA
 
-    client = ChatNVIDIA(**llm_config.model_dump(exclude={"type"}, by_alias=True))
-
-    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)
+    # prefer max_completion_tokens over max_tokens
+    client = ChatNVIDIA(
+        **llm_config.model_dump(exclude={"type", "max_tokens"}, by_alias=True),
+        max_completion_tokens=llm_config.max_tokens,
+    )
 
-    yield client
+    yield _patch_llm_based_on_config(client, llm_config)
 
 
 @register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
@@ -77,18 +144,7 @@ async def openai_langchain(llm_config: OpenAIModelConfig, _builder: Builder):
 
     from langchain_openai import ChatOpenAI
 
-    # Default kwargs for OpenAI to include usage metadata in the response. If the user has set stream_usage to False, we
-    # will not include this.
-    default_kwargs = {"stream_usage": True}
-
-    kwargs = {**default_kwargs, **llm_config.model_dump(exclude={"type"}, by_alias=True)}
-
-    client = ChatOpenAI(**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 stream_usage is specified, it will override the default value of True.
+    client = ChatOpenAI(stream_usage=True, **llm_config.model_dump(exclude={"type"}, by_alias=True))
 
-    yield client
+    yield _patch_llm_based_on_config(client, llm_config)

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

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-langchain
-Version: 1.3.0a20250826
+Version: 1.3.0a20250828
 Summary: Subpackage for LangChain and LangGraph 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: langchain-aws~=0.2.1
 Requires-Dist: langchain-core~=0.3.7
 Requires-Dist: langchain-nvidia-ai-endpoints~=0.3.5

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

@@ -1,7 +1,7 @@
 nat/meta/pypi.md,sha256=-RewrXPwhrT6398iluvXb5lefn18PybmvRFhmZF7KVI,1124
 nat/plugins/langchain/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nat/plugins/langchain/embedder.py,sha256=7bHXRcLwCeqG6ZQZQ35UAb46k5bN3UUL02OtG7RdL4Y,3282
-nat/plugins/langchain/llm.py,sha256=HJZsk0i72ZN0QTb1hZSUV4IsLmE0Ab0JOYYnvtqJfno,4013
+nat/plugins/langchain/llm.py,sha256=VbBq5wuTz8VjQjiUVE-bBfsQr2Zkx5wwNiTbvi-Ngaw,6537
 nat/plugins/langchain/register.py,sha256=jgq6wSJoGQIZFJhS8RbUs25cLgNJjCkFu4M6qaWJS_4,906
 nat/plugins/langchain/retriever.py,sha256=SWbXXOezEUuPACnmSSU497NAmEVEMj2SrFJGodkRg34,2644
 nat/plugins/langchain/tool_wrapper.py,sha256=Zgb2_XB4bEhjPPeqS-ZH_OJT_pcQmteX7u03N_qCLfc,2121
@@ -10,8 +10,8 @@ nat/plugins/langchain/tools/code_generation_tool.py,sha256=qL3HBiOQzVPLw4EiUOWes
 nat/plugins/langchain/tools/register.py,sha256=uemxqLxcNk1bGX4crV52oMphLTZWonStzkXwTZeG2Rw,889
 nat/plugins/langchain/tools/tavily_internet_search.py,sha256=AnnLRY1xSU4DOzxbB8nFZRjHngXpqatPVOJ7yWV7jVw,2612
 nat/plugins/langchain/tools/wikipedia_search.py,sha256=431YwLsjoC_mdvMZ_gY0Q37Uqaue2ASnAHpwr4jWCaU,2197
-nvidia_nat_langchain-1.3.0a20250826.dist-info/METADATA,sha256=ZOIis3NGacByh78Sqvro8zVYc6qCus0BEGqxQ8yfNSI,1735
-nvidia_nat_langchain-1.3.0a20250826.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_langchain-1.3.0a20250826.dist-info/entry_points.txt,sha256=4deXsMn97I012HhDw0UjoqcZ8eEoZ7BnqaRx5QmzebY,123
-nvidia_nat_langchain-1.3.0a20250826.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_langchain-1.3.0a20250826.dist-info/RECORD,,
+nvidia_nat_langchain-1.3.0a20250828.dist-info/METADATA,sha256=hBxPAKOpltN-iXiHXlSr49pnRNtEgS8sngYGYxuLE8I,1735
+nvidia_nat_langchain-1.3.0a20250828.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_langchain-1.3.0a20250828.dist-info/entry_points.txt,sha256=4deXsMn97I012HhDw0UjoqcZ8eEoZ7BnqaRx5QmzebY,123
+nvidia_nat_langchain-1.3.0a20250828.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_langchain-1.3.0a20250828.dist-info/RECORD,,