feat(llm): add provider-agnostic parameter mapping system

Pouyanpi · Pouyanpi · commit ee4e724eb9be · 2025-10-02T15:34:46.000+02:00
Implements flexible LLM parameter transformation to support
provider-specific naming conventions (e.g., max_tokens -&gt; max_new_tokens
for HuggingFace).
diff --git a/nemoguardrails/actions/llm/utils.py b/nemoguardrails/actions/llm/utils.py
@@ -30,6 +30,7 @@
     tool_calls_var,
 )
 from nemoguardrails.integrations.langchain.message_utils import dicts_to_messages
+from nemoguardrails.llm.parameter_mapping import transform_llm_params
 from nemoguardrails.logging.callbacks import logging_callbacks
 from nemoguardrails.logging.explain import LLMCallInfo
 
@@ -97,9 +98,18 @@ async def llm_call(
     _setup_llm_call_info(llm, model_name, model_provider)
     all_callbacks = _prepare_callbacks(custom_callback_handlers)
 
-    generation_llm: Union[BaseLanguageModel, Runnable] = (
-        llm.bind(stop=stop, **llm_params) if llm_params and llm is not None else llm
-    )
+    if llm_params or stop:
+        params_to_transform = llm_params.copy() if llm_params else {}
+        if stop is not None:
+            params_to_transform["stop"] = stop
+        transformed_params = transform_llm_params(
+            params_to_transform, llm, model_provider
+        )
+        generation_llm: Union[BaseLanguageModel, Runnable] = llm.bind(
+            **transformed_params
+        )
+    else:
+        generation_llm: Union[BaseLanguageModel, Runnable] = llm
 
     if isinstance(prompt, str):
         response = await _invoke_with_string_prompt(
diff --git a/nemoguardrails/llm/parameter_mapping.py b/nemoguardrails/llm/parameter_mapping.py
@@ -0,0 +1,171 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Module for transforming LLM parameters between internal and provider-specific formats."""
+
+import logging
+import weakref
+from typing import Any, Dict, Optional
+
+from langchain.base_language import BaseLanguageModel
+
+log = logging.getLogger(__name__)
+
+# Global registry to store parameter mappings for LLM instances
+_llm_parameter_mappings = weakref.WeakKeyDictionary()
+
+PROVIDER_PARAMETER_MAPPINGS = {
+    "huggingface": {
+        "max_tokens": "max_new_tokens",
+    },
+    "google_vertexai": {
+        "max_tokens": "max_output_tokens",
+    },
+}
+
+
+def register_llm_parameter_mapping(
+    llm: BaseLanguageModel, parameter_mapping: Dict[str, Optional[str]]
+) -> None:
+    """Register a parameter mapping for a specific LLM instance.
+
+    Args:
+        llm: The LLM instance
+        parameter_mapping: The parameter mapping dictionary
+    """
+    _llm_parameter_mappings[llm] = parameter_mapping
+    log.debug(f"Registered parameter mapping for LLM {type(llm).__name__}")
+
+
+def get_llm_parameter_mapping(
+    llm: BaseLanguageModel,
+) -> Optional[Dict[str, Optional[str]]]:
+    """Get the registered parameter mapping for an LLM instance.
+
+    Args:
+        llm: The LLM instance
+
+    Returns:
+        The parameter mapping if registered, None otherwise
+    """
+    return _llm_parameter_mappings.get(llm)
+
+
+def _infer_provider_from_module(llm: BaseLanguageModel) -> Optional[str]:
+    """Infer provider name from the LLM's module path.
+
+    This function extracts the provider name from LangChain package naming conventions:
+    - langchain_openai -> openai
+    - langchain_anthropic -> anthropic
+    - langchain_google_genai -> google_genai
+    - langchain_nvidia_ai_endpoints -> nvidia_ai_endpoints
+    - langchain_community.chat_models.ollama -> ollama
+
+    Args:
+        llm: The LLM instance
+
+    Returns:
+        The inferred provider name, or None if it cannot be determined
+    """
+    module = type(llm).__module__
+
+    if module.startswith("langchain_"):
+        package = module.split(".")[0]
+        provider = package.replace("langchain_", "")
+
+        if provider == "community":
+            parts = module.split(".")
+            if len(parts) >= 3:
+                provider = parts[-1]
+                log.debug(
+                    "Inferred provider '%s' from community module %s", provider, module
+                )
+                return provider
+        else:
+            log.debug("Inferred provider '%s' from module %s", provider, module)
+            return provider
+
+    log.debug(f"Could not infer provider from module {module}")
+    return None
+
+
+def get_llm_provider(llm: BaseLanguageModel) -> Optional[str]:
+    """Get the provider name for an LLM instance by inferring from module path.
+
+    This function extracts the provider name from LangChain package naming conventions.
+    See _infer_provider_from_module for details on the inference logic.
+
+    Args:
+        llm: The LLM instance
+
+    Returns:
+        The provider name if it can be inferred, None otherwise
+    """
+    return _infer_provider_from_module(llm)
+
+
+def transform_llm_params(
+    llm_params: Dict[str, Any],
+    llm: BaseLanguageModel,
+    provider: Optional[str] = None,
+    parameter_mapping: Optional[Dict[str, Optional[str]]] = None,
+) -> Dict[str, Any]:
+    """Transform LLM parameters using provider-specific or custom mappings.
+
+    Args:
+        llm_params: The original parameters dictionary
+        llm: The LLM instance to infer provider from
+        provider: Optional provider name. If None, will be automatically determined from llm.
+        parameter_mapping: Custom mapping dictionary. If None, uses built-in provider mappings.
+                          Key is the internal parameter name, value is the provider parameter name.
+                          If value is None, the parameter is dropped.
+
+    Returns:
+        Transformed parameters dictionary
+    """
+    if not llm_params:
+        return llm_params
+
+    mapping = parameter_mapping
+    if mapping is None:
+        mapping = get_llm_parameter_mapping(llm)
+        if mapping:
+            log.debug("Using registered parameter mapping for LLM instance")
+        else:
+            if provider is None:
+                provider = get_llm_provider(llm)
+
+            if provider and provider in PROVIDER_PARAMETER_MAPPINGS:
+                mapping = PROVIDER_PARAMETER_MAPPINGS[provider]
+                log.debug("Using built-in parameter mapping for provider: %s", provider)
+            else:
+                return llm_params
+
+    if not mapping:
+        return llm_params
+
+    transformed_params = {}
+
+    for param_name, param_value in llm_params.items():
+        if param_name in mapping:
+            mapped_name = mapping[param_name]
+            if mapped_name is not None:
+                log.debug("Mapped parameter %s -> %s", param_name, mapped_name)
+            else:
+                log.debug("Dropped parameter %s", param_name)
+        else:
+            transformed_params[param_name] = param_value
+
+    return transformed_params
diff --git a/nemoguardrails/rails/llm/config.py b/nemoguardrails/rails/llm/config.py
@@ -123,6 +123,12 @@ class Model(BaseModel):
         description="Configuration parameters for reasoning LLMs.",
     )
     parameters: Dict[str, Any] = Field(default_factory=dict)
+    parameter_mapping: Optional[Dict[str, Optional[str]]] = Field(
+        default=None,
+        description="Optional parameter mapping to transform parameter names for provider-specific requirements. "
+        "Keys are internal parameter names, values are provider parameter names. "
+        "Set value to null to drop a parameter.",
+    )
 
     mode: Literal["chat", "text"] = Field(
         default="chat",
diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py
@@ -74,6 +74,7 @@
     ModelInitializationError,
     init_llm_model,
 )
+from nemoguardrails.llm.parameter_mapping import register_llm_parameter_mapping
 from nemoguardrails.logging.explain import ExplainInfo
 from nemoguardrails.logging.processing_log import compute_generation_log
 from nemoguardrails.logging.stats import LLMStats
@@ -443,11 +444,19 @@ def _init_llms(self):
         if self.llm:
             # If an LLM was provided via constructor, use it as the main LLM
             # Log a warning if a main LLM is also specified in the config
-            if any(model.type == "main" for model in self.config.models):
+            main_model = next(
+                (model for model in self.config.models if model.type == "main"), None
+            )
+            if main_model:
                 log.warning(
                     "Both an LLM was provided via constructor and a main LLM is specified in the config. "
                     "The LLM provided via constructor will be used and the main LLM from config will be ignored."
                 )
+                # Still register parameter mapping from config if available
+                if main_model.parameter_mapping:
+                    register_llm_parameter_mapping(
+                        self.llm, main_model.parameter_mapping
+                    )
             self.runtime.register_action_param("llm", self.llm)
 
             self._configure_main_llm_streaming(self.llm)
@@ -465,6 +474,10 @@ def _init_llms(self):
                     mode="chat",
                     kwargs=kwargs,
                 )
+                if main_model.parameter_mapping:
+                    register_llm_parameter_mapping(
+                        self.llm, main_model.parameter_mapping
+                    )
                 self.runtime.register_action_param("llm", self.llm)
 
                 self._configure_main_llm_streaming(
@@ -500,6 +513,10 @@ def _init_llms(self):
                     kwargs=kwargs,
                 )
 
+                if llm_config.parameter_mapping:
+                    register_llm_parameter_mapping(
+                        llm_model, llm_config.parameter_mapping
+                    )
                 if llm_config.type == "main":
                     # If a main LLM was already injected, skip creating another
                     # one. Otherwise, create and register it.
diff --git a/tests/test_llm_call_parameter_mapping.py b/tests/test_llm_call_parameter_mapping.py
@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Tests for LLM parameter mapping integration in llm_call function."""
+
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+
+from nemoguardrails.actions.llm.utils import llm_call
+
+
+class MockResponse:
+    """Mock response object."""
+
+    def __init__(self, content="Test response"):
+        self.content = content
+
+
+class MockHuggingFaceLLM:
+    """Mock HuggingFace LLM for testing parameter mapping."""
+
+    __module__ = "langchain_huggingface.llms"
+
+    def __init__(self):
+        self.bind = Mock(return_value=self)
+        self.ainvoke = AsyncMock(return_value=MockResponse())
+
+
+@pytest.mark.asyncio
+async def test_llm_call_with_registered_parameter_mapping():
+    """Test llm_call applies registered parameter mapping correctly."""
+    from nemoguardrails.llm.parameter_mapping import register_llm_parameter_mapping
+
+    mock_llm = MockHuggingFaceLLM()
+    register_llm_parameter_mapping(mock_llm, {"max_tokens": "max_new_tokens"})
+
+    result = await llm_call(
+        llm=mock_llm,
+        prompt="Test prompt",
+        llm_params={"max_tokens": 100, "temperature": 0.5},
+    )
+
+    mock_llm.bind.assert_called_once_with(max_new_tokens=100, temperature=0.5)
+    assert result == "Test response"
+
+
+@pytest.mark.asyncio
+async def test_llm_call_with_builtin_mapping():
+    """Test llm_call uses built-in provider mapping when no custom mapping provided."""
+    mock_llm = MockHuggingFaceLLM()
+
+    result = await llm_call(
+        llm=mock_llm,
+        prompt="Test prompt",
+        llm_params={"max_tokens": 50, "temperature": 0.7},
+    )
+
+    mock_llm.bind.assert_called_once_with(max_new_tokens=50, temperature=0.7)
+    assert result == "Test response"
+
+
+@pytest.mark.asyncio
+async def test_llm_call_with_dropped_parameter():
+    """Test llm_call drops parameters mapped to None."""
+    from nemoguardrails.llm.parameter_mapping import register_llm_parameter_mapping
+
+    mock_llm = MockHuggingFaceLLM()
+    register_llm_parameter_mapping(
+        mock_llm, {"max_tokens": "max_new_tokens", "unsupported_param": None}
+    )
+
+    result = await llm_call(
+        llm=mock_llm,
+        prompt="Test prompt",
+        llm_params={"max_tokens": 100, "unsupported_param": "value"},
+    )
+
+    mock_llm.bind.assert_called_once_with(max_new_tokens=100)
+    assert result == "Test response"
+
+
+@pytest.mark.asyncio
+async def test_llm_call_without_params():
+    """Test llm_call works without llm_params."""
+    mock_llm = MockHuggingFaceLLM()
+
+    result = await llm_call(llm=mock_llm, prompt="Test prompt")
+
+    mock_llm.bind.assert_not_called()
+    mock_llm.ainvoke.assert_called_once()
+    assert result == "Test response"
+
+
+@pytest.mark.asyncio
+async def test_llm_call_with_stop_tokens():
+    """Test llm_call handles stop tokens correctly with parameter mapping."""
+    from nemoguardrails.llm.parameter_mapping import register_llm_parameter_mapping
+
+    mock_llm = MockHuggingFaceLLM()
+    register_llm_parameter_mapping(mock_llm, {"max_tokens": "max_new_tokens"})
+
+    result = await llm_call(
+        llm=mock_llm,
+        prompt="Test prompt",
+        stop=["END", "STOP"],
+        llm_params={"max_tokens": 100},
+    )
+
+    mock_llm.bind.assert_called_once_with(stop=["END", "STOP"], max_new_tokens=100)
+    assert result == "Test response"
diff --git a/tests/test_llmrails.py b/tests/test_llmrails.py
diff --git a/tests/test_parameter_mapping.py b/tests/test_parameter_mapping.py
