Clean up failing tests, add new ones with api key checks in kwargs

Author: tgasser-nv

nemoguardrails/llm/models/initializer.py

@@ -23,12 +23,11 @@
 from .langchain_initializer import ModelInitializationError, init_langchain_model
 
 
-# later we can easily convert it to a class
+# later we can easily conver it to a class
 def init_llm_model(
     model_name: Optional[str],
     provider_name: str,
     mode: Literal["chat", "text"],
-    api_key: str,
     kwargs: Dict[str, Any],
 ) -> Union[BaseChatModel, BaseLLM]:
     """Initialize an LLM model with proper error handling.
@@ -40,7 +39,6 @@ def init_llm_model(
         model_name: Name of the model to initialize
         provider_name: Name of the provider to use
         mode: Literal taking either "chat" or "text" values
-        api_key: String with LLM API Key to use by client
         kwargs: Additional arguments to pass to the model initialization
 
     Returns:
@@ -54,7 +52,6 @@ def init_llm_model(
         model_name=model_name,
         provider_name=provider_name,
         mode=mode,
-        api_key=api_key,
         kwargs=kwargs,
     )
 

nemoguardrails/llm/models/langchain_initializer.py

@@ -47,7 +47,7 @@ class ModelInitializationError(Exception):
 
 
 ModelInitMethod = Callable[
-    [str, str, str, Dict[str, Any]], Optional[Union[BaseChatModel, BaseLLM]]
+    [str, str, Dict[str, Any]], Optional[Union[BaseChatModel, BaseLLM]]
 ]
 
 
@@ -67,10 +67,10 @@ def supports_mode(self, mode: Literal["chat", "text"]) -> bool:
         return mode in self.supported_modes
 
     def execute(
-        self, model_name: str, provider_name: str, api_key: str, kwargs: Dict[str, Any]
+        self, model_name: str, provider_name: str, kwargs: Dict[str, Any]
     ) -> Optional[Union[BaseChatModel, BaseLLM]]:
         """Execute this initializer to initialize a model."""
-        return self.init_method(model_name, provider_name, api_key, kwargs)
+        return self.init_method(model_name, provider_name, kwargs)
 
     def __str__(self) -> str:
         return f"{self.init_method.__name__}(modes={self.supported_modes})"
@@ -81,7 +81,6 @@ def try_initialization_method(
     model_name: str,
     provider_name: str,
     mode: Literal["chat", "text"],
-    api_key: str,
     kwargs: Dict[str, Any],
 ):
     """Wrap an initialization method execution with a try/except to capture errors.
@@ -106,7 +105,6 @@ def try_initialization_method(
         result = initializer.execute(
             model_name=model_name,
             provider_name=provider_name,
-            api_key=api_key,
             kwargs=kwargs,
         )
         log.debug(f"Initializer {initializer.init_method.__name__} returned: {result}")
@@ -124,7 +122,6 @@ def init_langchain_model(
     model_name: str,
     provider_name: str,
     mode: Literal["chat", "text"],
-    api_key: str,
     kwargs: Dict[str, Any],
 ) -> Union[BaseChatModel, BaseLLM]:
     """Initialize a LangChain model using a series of initialization methods.
@@ -165,7 +162,6 @@ def init_langchain_model(
                 model_name=model_name,
                 provider_name=provider_name,
                 mode=mode,
-                api_key=api_key,
                 kwargs=kwargs,
             )
             if result is not None:
@@ -201,14 +197,13 @@ def init_langchain_model(
 
 
 def _init_chat_completion_model(
-    model_name: str, provider_name: str, api_key: str, kwargs: Dict[str, Any]
+    model_name: str, provider_name: str, kwargs: Dict[str, Any]
 ) -> BaseChatModel:  # noqa #type: ignore
     """Initialize a chat completion model.
 
     Args:
         model_name: Name of the model to initialize
         provider_name: Name of the provider to use
-        api_key: LLM API key to initialize client with
         kwargs: Additional arguments to pass to the model initialization
 
     Returns:
@@ -223,10 +218,6 @@ def _init_chat_completion_model(
     # line with our pyproject.toml
     package_version = version("langchain-core")
 
-    # Langchain's `init_chat_model()` doesn't have an argument for api_key, so
-    # copy kwargs and include api_key there instead
-    kwargs["api_key"] = api_key
-
     if _parse_version(package_version) < (0, 2, 7):
         raise RuntimeError(
             "this feature is supported from v0.2.7 of langchain-core."
@@ -243,14 +234,13 @@ def _init_chat_completion_model(
 
 
 def _init_text_completion_model(
-    model_name: str, provider_name: str, api_key: str, kwargs: Dict[str, Any]
+    model_name: str, provider_name: str, kwargs: Dict[str, Any]
 ) -> BaseLLM:
     """Initialize a text completion model.
 
     Args:
         model_name: Name of the model to initialize
         provider_name: Name of the provider to use
-        api_key: API Key to use for LLM call
         kwargs: Additional arguments to pass to the model initialization
 
     Returns:
@@ -262,20 +252,18 @@ def _init_text_completion_model(
     provider_cls = _get_text_completion_provider(provider_name)
     if provider_cls is None:
         raise ValueError()
-    kwargs = _update_model_kwargs(provider_cls, model_name, api_key, kwargs)
-    kwargs["api_key"] = api_key
+    kwargs = _update_model_kwargs(provider_cls, model_name, kwargs)
     return provider_cls(**kwargs)
 
 
 def _init_community_chat_models(
-    model_name: str, provider_name: str, api_key: str, kwargs: Dict[str, Any]
+    model_name: str, provider_name: str, kwargs: Dict[str, Any]
 ) -> BaseChatModel:
     """Initialize community chat models.
 
     Args:
         provider_name: Name of the provider to use
         model_name: Name of the model to initialize
-        api_key: API Key to use for LLM call
         kwargs: Additional arguments to pass to the model initialization
 
     Returns:
@@ -288,12 +276,12 @@ def _init_community_chat_models(
     provider_cls = _get_chat_completion_provider(provider_name)
     if provider_cls is None:
         raise ValueError()
-    kwargs = _update_model_kwargs(provider_cls, model_name, api_key, kwargs)
+    kwargs = _update_model_kwargs(provider_cls, model_name, kwargs)
     return provider_cls(**kwargs)
 
 
 def _init_gpt35_turbo_instruct(
-    model_name: str, provider_name: str, api_key: str, kwargs: Dict[str, Any]
+    model_name: str, provider_name: str, kwargs: Dict[str, Any]
 ) -> BaseLLM:
     """Initialize GPT-3.5 Turbo Instruct model.
 
@@ -305,7 +293,6 @@ def _init_gpt35_turbo_instruct(
     Args:
         model_name: Name of the model to initialize
         provider_name: Name of the provider to use
-        api_key: API key value for LLM call
         kwargs: Additional arguments to pass to the model initialization
 
     Returns:
@@ -318,7 +305,6 @@ def _init_gpt35_turbo_instruct(
         return _init_text_completion_model(
             model_name=model_name,
             provider_name=provider_name,
-            api_key=api_key,
             kwargs=kwargs,
         )
     except Exception as e:
@@ -328,14 +314,13 @@ def _init_gpt35_turbo_instruct(
 
 
 def _init_nvidia_model(
-    model_name: str, provider_name: str, api_key: str, kwargs: Dict[str, Any]
+    model_name: str, provider_name: str, kwargs: Dict[str, Any]
 ) -> BaseChatModel:
     """Initialize NVIDIA AI Endpoints model.
 
     Args:
         model_name: Name of the model to initialize
         provider_name: Name of the provider to use
-        api_key: API key
         **kwargs: Additional arguments to pass to the model initialization
 
     Returns:
@@ -358,7 +343,7 @@ def _init_nvidia_model(
                 " Please upgrade it with `pip install langchain-nvidia-ai-endpoints --upgrade`."
             )
 
-        return ChatNVIDIA(model=model_name, api_key=api_key, **kwargs)
+        return ChatNVIDIA(model=model_name, **kwargs)
     except ImportError as e:
         raise ImportError(
             "Could not import langchain_nvidia_ai_endpoints, please install it with "
@@ -379,7 +364,7 @@ def _init_nvidia_model(
 
 
 def _handle_model_special_cases(
-    model_name: str, provider_name: str, api_key: str, kwargs: Dict[str, Any]
+    model_name: str, provider_name: str, kwargs: Dict[str, Any]
 ) -> Optional[Union[BaseChatModel, BaseLLM]]:
     """Handle model initialization for special cases that need custom logic.
 
@@ -408,15 +393,13 @@ def _handle_model_special_cases(
     if initializer is None:
         return None
 
-    result = initializer(model_name, provider_name, api_key, kwargs)
+    result = initializer(model_name, provider_name, kwargs)
     if not isinstance(result, (BaseChatModel, BaseLLM)):
         raise TypeError("Initializer returned an invalid type")
     return result
 
 
-def _update_model_kwargs(
-    provider_cls: type, model_name: str, api_key: str, kwargs: dict
-) -> Dict:
+def _update_model_kwargs(provider_cls: type, model_name: str, kwargs: dict) -> Dict:
     """Update kwargs with the model name based on the provider's expected fields.
 
     If provider_cls.model_fields contains 'model' or 'model_name',
@@ -425,5 +408,4 @@ def _update_model_kwargs(
     for key in ("model", "model_name"):
         if key in getattr(provider_cls, "model_fields", {}):
             kwargs[key] = model_name
-    kwargs["api_key"] = api_key
     return kwargs

nemoguardrails/llm/providers/providers.py

@@ -133,7 +133,7 @@ def _patch_acall_method_to(llm_providers: Dict[str, Type[BaseLLM]]):
             setattr(provider_cls, "_acall", _acall)
 
 
-# Initialize the providers with the default onesBTW
+# Initialize the providers with the default ones
 _llm_providers: Dict[str, Type[BaseLLM]] = {
     "trt_llm": TRTLLM,
 }

nemoguardrails/rails/llm/config.py

@@ -102,9 +102,9 @@ class Model(BaseModel):
         default=None,
         description="The name of the model. If not specified, it should be specified through the parameters attribute.",
     )
-    api_key_env_var: str = Field(
+    api_key_env_var: Optional[str] = Field(
         default=None,
-        description='The environment variable containing the model\'s API Key. Do not include "$".',
+        description='Optional environment variable with model\'s API Key. Do not include "$".',
     )
     reasoning_config: Optional[ReasoningModelConfig] = Field(
         default_factory=ReasoningModelConfig,
@@ -1245,7 +1245,7 @@ def validate_models_api_key_env_var(cls, models):
         """Model API Key Env var must be set to make LLM calls"""
         api_keys = [m.api_key_env_var for m in models]
         for api_key in api_keys:
-            if not os.environ.get(api_key):
+            if api_key and not os.environ.get(api_key):
                 raise ValueError(
                     f"Model API Key environment variable '{api_key}' not set."
                 )

nemoguardrails/rails/llm/llmrails.py

@@ -376,18 +376,18 @@ def _init_llms(self):
                 provider_name = llm_config.engine
                 kwargs = llm_config.parameters or {}
                 mode = llm_config.mode
-                api_key = os.environ.get(llm_config.api_key_env_var)
 
-                # Add the api_key to kwargs if it's set
-                api_key = os.environ.get(llm_config.api_key_env_var)
-                if api_key:
-                    kwargs["api_key"] = api_key
+                # If the optional API Key Environment Variable is set, store
+                # this in the `kwargs` for the current model
+                if llm_config.api_key_env_var:
+                    api_key = os.environ.get(llm_config.api_key_env_var)
+                    if api_key:
+                        kwargs["api_key"] = api_key
 
                 llm_model = init_llm_model(
                     model_name=model_name,
                     provider_name=provider_name,
                     mode=mode,
-                    api_key=api_key,
                     kwargs=kwargs,
                 )