refactor: update NVIDIA System-Assist client initialization and error handling (#8278)

* refactor: Improve NVIDIA System-Assist client initialization and error handling

- Updated the `maybe_register_rise_client` method to enhance error handling and remove unnecessary context management.
- Clarified the component description to specify Windows-only support.
- Improved input prompt information with an example for better user guidance.

* refactor: use constant for rise_initialized key in NvidiaSystemAssistComponent

Replaced hardcoded string with a constant for the rise_initialized key to improve code readability and maintainability.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2025-06-02 17:38:39 -03:00 committed by GitHub
commit 4a0279f56f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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()