NVIDIA · Pouyanpi · Apr 17, 2025 · Apr 2, 2025 · Apr 16, 2025 · Apr 16, 2025
diff --git a/examples/configs/content_safety/config.yml b/examples/configs/content_safety/config.yml
@@ -1,7 +1,7 @@
 models:
   - type: main
     engine: nim
-    model_name: meta/llama-3.3-70b-instruct
+    model: meta/llama-3.3-70b-instruct
 
   - type: content_safety
     engine: nim

diff --git a/nemoguardrails/actions/llm/generation.py b/nemoguardrails/actions/llm/generation.py
@@ -23,10 +23,11 @@
 import threading
 from functools import lru_cache
 from time import time
-from typing import Callable, List, Optional, cast
+from typing import Callable, List, Optional, Union, cast
 
 from jinja2 import meta
 from jinja2.sandbox import SandboxedEnvironment
+from langchain_core.language_models import BaseChatModel
 from langchain_core.language_models.llms import BaseLLM
 
 from nemoguardrails.actions.actions import ActionResult, action
@@ -81,7 +82,7 @@ class LLMGenerationActions:
     def __init__(
         self,
         config: RailsConfig,
-        llm: BaseLLM,
+        llm: Union[BaseLLM, BaseChatModel],
         llm_task_manager: LLMTaskManager,
         get_embedding_search_provider_instance: Callable[
             [Optional[EmbeddingSearchProvider]], EmbeddingsIndex
@@ -417,7 +418,7 @@ async def generate_user_intent(
                     )
                 # We add these in reverse order so the most relevant is towards the end.
                 for result in reversed(results):
-                    examples += f"user \"{result.text}\"\n  {result.meta['intent']}\n\n"
+                    examples += f'user "{result.text}"\n  {result.meta["intent"]}\n\n'
                     if result.meta["intent"] not in potential_user_intents:
                         potential_user_intents.append(result.meta["intent"])
 

diff --git a/nemoguardrails/evaluate/utils.py b/nemoguardrails/evaluate/utils.py
@@ -15,27 +15,18 @@
 
 import json
 
-from nemoguardrails.llm.providers import get_llm_provider, get_llm_provider_names
+from nemoguardrails.llm.models.initializer import init_llm_model
 from nemoguardrails.rails.llm.config import Model
 
 
 def initialize_llm(model_config: Model):
     """Initializes the model from LLM provider."""
-    if model_config.engine not in get_llm_provider_names():
-        raise Exception(f"Unknown LLM engine: {model_config.engine}")
-    provider_cls = get_llm_provider(model_config)
-    kwargs = {"temperature": 0, "max_tokens": 10}
-    if model_config.engine in [
-        "azure",
-        "openai",
-        "gooseai",
-        "nlpcloud",
-        "petals",
-    ]:
-        kwargs["model_name"] = model_config.model
-    else:
-        kwargs["model"] = model_config.model
-    return provider_cls(**kwargs)
+
+    return init_llm_model(
+        model_name=model_config.model,
+        provider_name=model_config.engine,
+        kwargs=model_config.parameters,
+    )
 
 
 def load_dataset(dataset_path: str):

diff --git a/nemoguardrails/llm/models/initializer.py b/nemoguardrails/llm/models/initializer.py
@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module for initializing LLM models with proper error handling and type checking."""
+
+from typing import Any, Dict, Literal, Optional, Union
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.language_models.llms import BaseLLM
+
+from .langchain_initializer import ModelInitializationError, init_langchain_model
+
+
+# later we can easily conver it to a class
+def init_llm_model(
+    model_name: Optional[str],
+    provider_name: str,
+    mode: Literal["chat", "text"],
+    kwargs: Dict[str, Any],
+) -> Union[BaseChatModel, BaseLLM]:
+    """Initialize an LLM model with proper error handling.
+
+    Currently, this function only supports LangChain models.
+    In the future, it may support other model backends.
+
+    Args:
+        model_name: Name of the model to initialize
+        provider_name: Name of the provider to use
+        kwargs: Additional arguments to pass to the model initialization
+
+    Returns:
+        An initialized LLM model
+
+    Raises:
+        ModelInitializationError: If model initialization fails
+    """
+    # currently we only support LangChain models
+    return init_langchain_model(
+        model_name=model_name, provider_name=provider_name, mode=mode, kwargs=kwargs
+    )
+
+
+__all__ = ["init_llm_model", "ModelInitializationError"]