diff --git a/src/backend/base/langflow/components/nvidia/system_assist.py b/src/backend/base/langflow/components/nvidia/system_assist.py index 8cf5779d5..d68848fd8 100644 --- a/src/backend/base/langflow/components/nvidia/system_assist.py +++ b/src/backend/base/langflow/components/nvidia/system_assist.py @@ -1,46 +1,29 @@ import asyncio -import contextlib from langflow.custom.custom_component.component_with_cache import ComponentWithCache from langflow.io import MessageTextInput, Output from langflow.schema import Message from langflow.services.cache.utils import CacheMiss +RISE_INITIALIZED_KEY = "rise_initialized" + class NvidiaSystemAssistComponent(ComponentWithCache): display_name = "NVIDIA System-Assist" description = ( - "Prompts NVIDIA System-Assist to interact with the NVIDIA GPU Driver. " + "(Windows only) Prompts NVIDIA System-Assist to interact with the NVIDIA GPU Driver. " "The user may query GPU specifications, state, and ask the NV-API to perform " "several GPU-editing acations. The prompt must be human-readable language." - "(Windows only)" ) documentation = "https://docs.langflow.org/components-custom-components" icon = "NVIDIA" rise_initialized = False - def maybe_register_rise_client(self): - with contextlib.suppress(ImportError): - from gassist.rise import register_rise_client - rise_initialized = self._shared_component_cache.get("rise_initialized") - if not isinstance(rise_initialized, CacheMiss) and rise_initialized: - return - self.log("Initializing Rise Client") - try: - register_rise_client() - self._shared_component_cache.set(key="rise_initialized", value=True) - except NameError as e: - msg = "NVIDIA System-Assist is Windows only and not supported on this platform" - raise ValueError(msg) from e - except Exception as e: - msg = f"An error occurred initializing NVIDIA System-Assist: {e}" - raise ValueError(msg) from e - inputs = [ MessageTextInput( name="prompt", display_name="System-Assist Prompt", - info="Enter a prompt for NVIDIA System-Assist to process.", + info="Enter a prompt for NVIDIA System-Assist to process. Example: 'What is my GPU?'", value="", tool_mode=True, ), @@ -50,9 +33,30 @@ class NvidiaSystemAssistComponent(ComponentWithCache): Output(display_name="Response", name="response", method="sys_assist_prompt"), ] + def maybe_register_rise_client(self): + try: + from gassist.rise import register_rise_client + + rise_initialized = self._shared_component_cache.get(RISE_INITIALIZED_KEY) + if not isinstance(rise_initialized, CacheMiss) and rise_initialized: + return + self.log("Initializing Rise Client") + + register_rise_client() + self._shared_component_cache.set(key=RISE_INITIALIZED_KEY, value=True) + except ImportError as e: + msg = "NVIDIA System-Assist is Windows only and not supported on this platform" + raise ValueError(msg) from e + except Exception as e: + msg = f"An error occurred initializing NVIDIA System-Assist: {e}" + raise ValueError(msg) from e + async def sys_assist_prompt(self) -> Message: - with contextlib.suppress(ImportError): + try: from gassist.rise import send_rise_command + except ImportError as e: + msg = "NVIDIA System-Assist is Windows only and not supported on this platform" + raise ValueError(msg) from e self.maybe_register_rise_client()