Revert langchain_initializer changes

Author: JashG

nemoguardrails/llm/models/langchain_initializer.py

@@ -227,7 +227,7 @@ def _init_chat_completion_model(model_name: str, provider_name: str, kwargs: Dic
 
 def _init_text_completion_model(
     model_name: str, provider_name: str, kwargs: Dict[str, Any]
-) -> BaseLLM | None:
+) -> BaseLLM:
     """Initialize a text completion model.
 
     Args:
@@ -241,14 +241,9 @@ def _init_text_completion_model(
     Raises:
         RuntimeError: If the provider is not found
     """
-    try:
-        provider_cls = _get_text_completion_provider(provider_name)
-    except RuntimeError as e:
-        return None
-
+    provider_cls = _get_text_completion_provider(provider_name)
     if provider_cls is None:
-        return None
-
+        raise ValueError()
     kwargs = _update_model_kwargs(provider_cls, model_name, kwargs)
     # remove stream_usage parameter as it's not supported by text completion APIs
     # (e.g., OpenAI's AsyncCompletions.create() doesn't accept this parameter)
@@ -258,7 +253,7 @@ def _init_text_completion_model(
 
 def _init_community_chat_models(
     model_name: str, provider_name: str, kwargs: Dict[str, Any]
-) -> BaseChatModel | None:
+) -> BaseChatModel:
     """Initialize community chat models.
 
     Args:
@@ -275,14 +270,14 @@ def _init_community_chat_models(
     """
     provider_cls = _get_chat_completion_provider(provider_name)
     if provider_cls is None:
-        return None
+        raise ValueError()
     kwargs = _update_model_kwargs(provider_cls, model_name, kwargs)
     return provider_cls(**kwargs)
 
 
 def _init_gpt35_turbo_instruct(
     model_name: str, provider_name: str, kwargs: Dict[str, Any]
-) -> BaseLLM | None:
+) -> BaseLLM:
     """Initialize GPT-3.5 Turbo Instruct model.
 
     Currently init_chat_model from langchain infers this as a chat model.

tests/llm_providers/test_langchain_initialization_methods.py

@@ -116,9 +116,8 @@ def test_init_community_chat_models_no_provider(self):
             "nemoguardrails.llm.models.langchain_initializer._get_chat_completion_provider"
         ) as mock_get_provider:
             mock_get_provider.return_value = None
-            assert (
-                _init_community_chat_models("community-model", "provider", {}) is None
-            )
+            with pytest.raises(ValueError):
+                _init_community_chat_models("community-model", "provider", {})
 
 
 class TestTextCompletionInitializer:
@@ -157,7 +156,8 @@ def test_init_text_completion_model_no_provider(self):
             "nemoguardrails.llm.models.langchain_initializer._get_text_completion_provider"
         ) as mock_get_provider:
             mock_get_provider.return_value = None
-            assert _init_text_completion_model("text-model", "provider", {}) is None
+            with pytest.raises(ValueError):
+                _init_text_completion_model("text-model", "provider", {})
 
 
 class TestUpdateModelKwargs:

tests/llm_providers/test_langchain_initializer.py

@@ -183,14 +183,14 @@ def test_text_completion_called_when_previous_fail(mock_initializers):
     mock_initializers["text"].assert_called_once()
 
 
-def test_text_mode_only_calls_text_initializers(mock_initializers):
-    """Test that text mode only tries initializers that support text mode."""
+def test_text_completion_supports_chat_mode(mock_initializers):
     mock_initializers["special"].return_value = None
+    mock_initializers["chat"].return_value = None
+    mock_initializers["community"].return_value = None
     mock_initializers["text"].return_value = "text_model"
-    result = init_langchain_model("text-model", "provider", "text", {})
+    result = init_langchain_model("text-model", "provider", "chat", {})
     assert result == "text_model"
     mock_initializers["special"].assert_called_once()
+    mock_initializers["chat"].assert_called_once()
+    mock_initializers["community"].assert_called_once()
     mock_initializers["text"].assert_called_once()
-    # Since text returns a value, chat and community are not called
-    mock_initializers["chat"].assert_not_called()
-    mock_initializers["community"].assert_not_called()

tests/test_actions_llm_utils.py

@@ -91,24 +91,6 @@ class MockPatchedNVIDIA(MockNVIDIAOriginal):
     __module__ = "nemoguardrails.llm.providers._langchain_nvidia_ai_endpoints_patch"
 
 
-class MockTRTLLM:
-    __module__ = "nemoguardrails.llm.providers.trtllm.llm"
-
-
-class MockAzureLLM:
-    __module__ = "langchain_openai.chat_models"
-
-
-class MockLLMWithClient:
-    __module__ = "langchain_openai.chat_models"
-
-    class _MockClient:
-        base_url = "https://custom.endpoint.com/v1"
-
-    def __init__(self):
-        self.client = self._MockClient()
-
-
 def test_infer_provider_openai():
     llm = MockOpenAILLM()
     provider = _infer_provider_from_module(llm)