Add support for NVIDIA NIM as an LLM provider

.env

@@ -1,4 +1,4 @@
-# LLM Provider: "openai" (default), "anthropic", or "ollama"
+# LLM Provider: "openai" (default), "anthropic", "nvidia", or "ollama"
 LLM_PROVIDER=openai
 
 # OpenAI Configuration
@@ -14,3 +14,8 @@ ANTHROPIC_MODEL=claude-sonnet-4-5
 # Ollama Configuration (local models)
 OLLAMA_MODEL=llama3
 OLLAMA_BASE_URL=http://localhost:11434
+
+# NVIDIA NIM Configuration
+NVIDIA_API_KEY=
+NVIDIA_MODEL=nvidia/llama-3.3-nemotron-super-49b-v1.5
+NVIDIA_BASE_URL=https://integrate.api.nvidia.com/v1

pyproject.toml

@@ -40,8 +40,9 @@ dependencies = [
 
 [project.optional-dependencies]
 anthropic = ["langchain-anthropic~=0.3.12"]
+nvidia = ["langchain-nvidia-ai-endpoints~=0.3.9"]
 ollama = ["langchain-ollama~=0.3.2"]
-all = ["langchain-anthropic~=0.3.12", "langchain-ollama~=0.3.2"]
+all = ["langchain-anthropic~=0.3.12", "langchain-nvidia-ai-endpoints~=0.3.9", "langchain-ollama~=0.3.2"]
 
 [project.urls]
 "Homepage" = "https://github.com/nasa-jpl/rosa"

src/rosa/rosa.py

@@ -28,6 +28,7 @@
 
 if TYPE_CHECKING:
     from langchain_anthropic import ChatAnthropic
+    from langchain_nvidia_ai_endpoints import ChatNVIDIA
     from langchain_ollama import ChatOllama
 
 from .prompts import RobotSystemPrompts, system_prompts
@@ -37,7 +38,7 @@
 
 # Tested providers for static analysis; BaseChatModel accepted at runtime.
 if TYPE_CHECKING:
-    ChatModel = Union[ChatOpenAI, AzureChatOpenAI, ChatAnthropic, ChatOllama]
+    ChatModel = Union[ChatOpenAI, AzureChatOpenAI, ChatAnthropic, ChatNVIDIA, ChatOllama]
 else:
     ChatModel = BaseChatModel
 
@@ -49,7 +50,7 @@ class ROSA:
     Args:
         ros_version (Literal[1, 2]): The version of ROS that the agent will interact with.
         llm (ChatModel): The language model to use for generating responses. Tested providers:
-            ChatOpenAI, AzureChatOpenAI, ChatAnthropic, and ChatOllama. Other BaseChatModel
+            ChatOpenAI, AzureChatOpenAI, ChatAnthropic, ChatNVIDIA, and ChatOllama. Other BaseChatModel
             subclasses that support tool calling may work but are not officially tested.
             Note: token usage tracking is only supported for ChatOpenAI and AzureChatOpenAI.
         tools (Optional[list]): A list of additional LangChain tool functions to use with the agent.

src/turtle_agent/scripts/llm.py

@@ -21,16 +21,17 @@
 def get_llm(streaming: bool = False):
     """A helper function to get the LLM instance.
 
-    Supports OpenAI (default), Anthropic and Ollama models.
+    Supports OpenAI (default), Anthropic, NVIDIA NIM and Ollama models.
     Set the LLM_PROVIDER env variable to switch between providers:
       - "openai" (default): uses OPENAI_API_KEY
       - "anthropic": uses ANTHROPIC_API_KEY
+      - "nvidia": uses NVIDIA_API_KEY (NIM API)
       - "ollama": uses local Ollama instance
     """
     dotenv.load_dotenv(dotenv.find_dotenv())
 
     provider = os.getenv("LLM_PROVIDER", "openai").lower().strip()
-    supported = ("openai", "anthropic", "ollama")
+    supported = ("openai", "anthropic", "nvidia", "ollama")
     if provider not in supported:
         raise ValueError(
             f"Unknown LLM_PROVIDER: '{provider}'. Must be one of: {', '.join(supported)}"
@@ -55,6 +56,19 @@ def get_llm(streaming: bool = False):
             model=os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-5"),
             streaming=streaming,
         )
+    elif provider == "nvidia":
+        try:
+            from langchain_nvidia_ai_endpoints import ChatNVIDIA
+        except ImportError as e:
+            raise ImportError(
+                "Install the project's NVIDIA extra with: pip install 'jpl-rosa[nvidia]'"
+            ) from e
+        llm = ChatNVIDIA(
+            api_key=get_env_variable("NVIDIA_API_KEY"),
+            model=os.getenv("NVIDIA_MODEL", "nvidia/llama-3.3-nemotron-super-49b-v1.5"),
+            base_url=os.getenv("NVIDIA_BASE_URL", "https://integrate.api.nvidia.com/v1"),
+            streaming=streaming,
+        )
     elif provider == "ollama":
         try:
             from langchain_ollama import ChatOllama
@@ -91,7 +105,7 @@ def get_env_variable(var_name: str) -> str:
     raise a ValueError, making it easier to debug configuration issues.
     """
     value = os.getenv(var_name)
-    if value is None:
+    if not value or not value.strip():
         msg = f"Environment variable {var_name} is not set."
         raise ValueError(msg)
-    return value
+    return value.strip()