feat(llm): Add custom HTTP headers support to ChatNVIDIA provider

Add custom HTTP headers support to the ChatNVIDIA class patch, enabling
users to pass custom headers
(authentication tokens, request IDs, billing information, etc.) with all
requests to NVIDIA AI
endpoints.

Implementation Approach

- Added custom_headers optional field to ChatNVIDIA class with Pydantic
v2 compatibility
- Implemented runtime method wrapping that intercepts _client.get_req()
and _client.get_req_stream()
to merge custom headers with existing headers
- Included automatic version detection to ensure compatibility with
langchain-nvidia-ai-endpoints >=
0.3.0, with clear error messages for older versions
- Works with both synchronous invoke() and streaming requests, fully
compatible with VLM (Vision
Language Models)

Author: Pouyanpi

nemoguardrails/llm/providers/_langchain_nvidia_ai_endpoints_patch.py

@@ -13,16 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import logging
 from functools import wraps
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 
 from langchain_core.callbacks.manager import CallbackManagerForLLMRun
 from langchain_core.language_models.chat_models import generate_from_stream
 from langchain_core.messages import BaseMessage
 from langchain_core.outputs import ChatResult
 from langchain_nvidia_ai_endpoints import ChatNVIDIA as ChatNVIDIAOriginal
-from pydantic.v1 import Field
+from pydantic import Field
 
 log = logging.getLogger(__name__)
 
@@ -55,6 +56,46 @@ class ChatNVIDIA(ChatNVIDIAOriginal):
     streaming: bool = Field(
         default=False, description="Whether to use streaming or not"
     )
+    custom_headers: Optional[Dict[str, str]] = Field(
+        default=None, description="Custom HTTP headers to send with requests"
+    )
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        if self.custom_headers:
+            if not hasattr(self._client, "get_req"):
+                raise RuntimeError(
+                    "custom_headers requires langchain-nvidia-ai-endpoints >= 0.2.1. "
+                    "Your version uses a nested client structure that is not supported. "
+                    "Please upgrade: pip install --upgrade langchain-nvidia-ai-endpoints>=0.3.0"
+                )
+
+            sig = inspect.signature(self._client.get_req)
+            if "extra_headers" not in sig.parameters:
+                raise RuntimeError(
+                    "custom_headers requires langchain-nvidia-ai-endpoints >= 0.3.0. "
+                    "Your version does not support the extra_headers parameter. "
+                    "Please upgrade: pip install --upgrade langchain-nvidia-ai-endpoints>=0.3.0"
+                )
+
+            self._wrap_client_methods()
+
+    def _wrap_client_methods(self):
+        original_get_req = self._client.get_req
+        original_get_req_stream = self._client.get_req_stream
+
+        def wrapped_get_req(payload: dict = {}, extra_headers: dict = {}):
+            merged_headers = {**extra_headers, **self.custom_headers}
+            return original_get_req(payload=payload, extra_headers=merged_headers)
+
+        def wrapped_get_req_stream(payload: dict, extra_headers: dict = {}):
+            merged_headers = {**extra_headers, **self.custom_headers}
+            return original_get_req_stream(
+                payload=payload, extra_headers=merged_headers
+            )
+
+        object.__setattr__(self._client, "get_req", wrapped_get_req)
+        object.__setattr__(self._client, "get_req_stream", wrapped_get_req_stream)
 
     @stream_decorator
     def _generate(