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

* 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,21 +13,22 @@
 # 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__)
+log = logging.getLogger(__name__)  # pragma: no cover
 
 
-def stream_decorator(func):
+def stream_decorator(func):  # pragma: no cover
     @wraps(func)
     def wrapper(
         self,
@@ -51,10 +52,52 @@ def wrapper(
 
 # NOTE: this needs to have the same name as the original class,
 #   otherwise, there's a check inside `langchain-nvidia-ai-endpoints` that will fail.
-class ChatNVIDIA(ChatNVIDIAOriginal):
+class ChatNVIDIA(ChatNVIDIAOriginal):  # pragma: no cover
     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:
+            custom_headers_error = (
+                "custom_headers requires langchain-nvidia-ai-endpoints >= 0.3.0. "
+                "Your version does not support the required client structure or "
+                "extra_headers parameter. Please upgrade: "
+                "pip install --upgrade langchain-nvidia-ai-endpoints>=0.3.0"
+            )
+            if not hasattr(self._client, "get_req"):
+                raise RuntimeError(custom_headers_error)
+
+            sig = inspect.signature(self._client.get_req)
+            if "extra_headers" not in sig.parameters:
+                raise RuntimeError(custom_headers_error)
+
+            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 = None, extra_headers: dict = None):
+            payload = payload or {}
+            extra_headers = extra_headers or {}
+            merged_headers = {**extra_headers, **self.custom_headers}
+            return original_get_req(payload=payload, extra_headers=merged_headers)
+
+        def wrapped_get_req_stream(payload: dict = None, extra_headers: dict = None):
+            payload = payload or {}
+            extra_headers = extra_headers or {}
+            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(