feat: NVIDIA Model Component - QoL and detailed thinking toggle for reasoning models (#7070)

* add toggle

* QoL changes for nvidia model comp

* Add detailed thinking prompt

* ruff

* makes detailed thinking a constant

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Edwin Jose <edwin.jose@datastax.com>
This commit is contained in:
Jordan Frazier 2025-03-18 09:59:33 -07:00 • committed by GitHub
commit 683a95948f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 80 additions and 22 deletions

View file

@ -16,6 +16,11 @@ from langflow.inputs.inputs import BoolInput, InputTypes, MultilineInput
from langflow.schema.message import Message
from langflow.template.field.base import Output
# Enabled detailed thinking for NVIDIA reasoning models.
#
# Models are trained with this exact string. Do not update.
DETAILED_THINKING_PREFIX = "detailed thinking on\n\n"
class LCModelComponent(Component):
display_name: str = "Model Name"
@ -162,6 +167,24 @@ class LCModelComponent(Component):
stream: bool,
input_value: str | Message,
system_message: str | None = None,
):
if getattr(self, "detailed_thinking", False):
system_message = DETAILED_THINKING_PREFIX + (system_message or "")
return self._get_chat_result(
runnable=runnable,
stream=stream,
input_value=input_value,
system_message=system_message,
)
def _get_chat_result(
self,
*,
runnable: LanguageModel,
stream: bool,
input_value: str | Message,
system_message: str | None = None,
):
messages: list[BaseMessage] = []
if not input_value and not system_message:

View file

@ -12,6 +12,14 @@ class NVIDIAModelComponent(LCModelComponent):
description = "Generates text using NVIDIA LLMs."
icon = "NVIDIA"
try:
from langchain_nvidia_ai_endpoints import ChatNVIDIA
all_models = ChatNVIDIA().get_available_models()
except ImportError as e:
msg = "Please install langchain-nvidia-ai-endpoints to use the NVIDIA model."
raise ImportError(msg) from e
inputs = [
*LCModelComponent._base_inputs,
IntInput(
@ -23,29 +31,35 @@ class NVIDIAModelComponent(LCModelComponent):
DropdownInput(
name="model_name",
display_name="Model Name",
info="The name of the NVIDIA model to use.",
advanced=False,
options=[],
refresh_button=True,
value=None,
options=[model.id for model in all_models],
combobox=True,
real_time_refresh=True,
),
BoolInput(
name="detailed_thinking",
display_name="Detailed Thinking",
info="If true, the model will return a detailed thought process. Only supported by reasoning models.",
value=False,
show=False,
),
BoolInput(
name="tool_model_enabled",
display_name="Enable Tool Models",
info="If enabled, only show models that support tool-calling.",
advanced=False,
value=False,
real_time_refresh=True,
),
MessageTextInput(
name="base_url",
display_name="NVIDIA Base URL",
value="https://integrate.api.nvidia.com/v1",
refresh_button=True,
info="The base URL of the NVIDIA API. Defaults to https://integrate.api.nvidia.com/v1.",
real_time_refresh=True,
),
BoolInput(
name="tool_model_enabled",
display_name="Enable Tool Models",
info=(
"Select if you want to use models that can work with tools. If yes, only those models will be shown."
),
advanced=False,
value=False,
real_time_refresh=True,
),
SecretStrInput(
name="api_key",
display_name="NVIDIA API Key",
@ -58,7 +72,7 @@ class NVIDIAModelComponent(LCModelComponent):
name="temperature",
display_name="Temperature",
value=0.1,
info="Run inference with this temperature. Must by in the closed interval [0.0, 1.0].",
info="Run inference with this temperature.",
range_spec=RangeSpec(min=0, max=1, step=0.01),
advanced=True,
),
@ -72,21 +86,42 @@ class NVIDIAModelComponent(LCModelComponent):
]
def get_models(self, tool_model_enabled: bool | None = None) -> list[str]:
build_model = self.build_model()
if tool_model_enabled:
tool_models = [model for model in build_model.get_available_models() if model.supports_tools]
return [model.id for model in tool_models]
return [model.id for model in build_model.available_models]
try:
from langchain_nvidia_ai_endpoints import ChatNVIDIA
except ImportError as e:
msg = "Please install langchain-nvidia-ai-endpoints to use the NVIDIA model."
raise ImportError(msg) from e
def update_build_config(self, build_config: dotdict, field_value: Any, field_name: str | None = None):
if field_name in {"base_url", "model_name", "tool_model_enabled", "api_key"} and field_value:
# Note: don't include the previous model, as it may not exist in available models from the new base url
model = ChatNVIDIA(base_url=self.base_url, api_key=self.api_key)
if tool_model_enabled:
tool_models = [m for m in model.get_available_models() if m.supports_tools]
return [m.id for m in tool_models]
return [m.id for m in model.available_models]
def update_build_config(self, build_config: dotdict, _field_value: Any, field_name: str | None = None):
if field_name in {"model_name", "tool_model_enabled", "base_url", "api_key"}:
try:
ids = self.get_models(self.tool_model_enabled)
build_config["model_name"]["options"] = ids
build_config["model_name"]["value"] = ids[0]
if "value" not in build_config["model_name"] or build_config["model_name"]["value"] is None:
build_config["model_name"]["value"] = ids[0]
elif build_config["model_name"]["value"] not in ids:
build_config["model_name"]["value"] = None
# TODO: use api to determine if model supports detailed thinking
if build_config["model_name"]["value"] == "nemotron":
build_config["detailed_thinking"]["show"] = True
else:
build_config["detailed_thinking"]["value"] = False
build_config["detailed_thinking"]["show"] = False
except Exception as e:
msg = f"Error getting model names: {e}"
build_config["model_name"]["value"] = None
build_config["model_name"]["options"] = []
raise ValueError(msg) from e
return build_config
def build_model(self) -> LanguageModel: # type: ignore[type-var]