Revert model initializer order, but only try _init_text_completion_model for text models

Author: JashG

nemoguardrails/llm/models/langchain_initializer.py

@@ -142,13 +142,12 @@ def init_langchain_model(
     initializers: list[ModelInitializer] = [
         # Try special case handlers first (handles both chat and text)
         ModelInitializer(_handle_model_special_cases, ["chat", "text"]),
-        # FIXME: is text and chat a good idea?
-        # For text mode, use text completion, we are using both text and chat as the last resort
-        ModelInitializer(_init_text_completion_model, ["text", "chat"]),
         # For chat mode, first try the standard chat completion API
         ModelInitializer(_init_chat_completion_model, ["chat"]),
         # For chat mode, fall back to community chat models
         ModelInitializer(_init_community_chat_models, ["chat"]),
+        # For text mode, use text completion
+        ModelInitializer(_init_text_completion_model, ["text"]),
     ]
 
     # Track the last exception for better error reporting

tests/llm_providers/test_langchain_initializer.py

@@ -70,27 +70,25 @@ def test_special_case_called_first(mock_initializers):
 
 def test_chat_completion_called(mock_initializers):
     mock_initializers["special"].return_value = None
-    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = "chat_model"
     result = init_langchain_model("chat-model", "provider", "chat", {})
     assert result == "chat_model"
     mock_initializers["special"].assert_called_once()
-    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_not_called()
+    mock_initializers["text"].assert_not_called()
 
 
 def test_community_chat_called(mock_initializers):
     mock_initializers["special"].return_value = None
-    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = None
     mock_initializers["community"].return_value = "community_model"
     result = init_langchain_model("community-chat", "provider", "chat", {})
     assert result == "community_model"
     mock_initializers["special"].assert_called_once()
-    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_called_once()
+    mock_initializers["text"].assert_not_called()
 
 
 def test_text_completion_called(mock_initializers):
@@ -110,13 +108,12 @@ def test_all_initializers_fail(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 = None
     with pytest.raises(ModelInitializationError):
         init_langchain_model("unknown-model", "provider", "chat", {})
     mock_initializers["special"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_called_once()
-    mock_initializers["text"].assert_called_once()
+    mock_initializers["text"].assert_not_called()
 
 
 def test_unsupported_mode(mock_initializers):
@@ -151,44 +148,41 @@ def test_all_initializers_raise_exceptions(mock_initializers):
     mock_initializers["special"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_called_once()
-    mock_initializers["text"].assert_called_once()
+    mock_initializers["text"].assert_not_called()
 
 
 def test_duplicate_modes_in_initializer(mock_initializers):
     mock_initializers["special"].return_value = None
-    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = "chat_model"
     result = init_langchain_model("chat-model", "provider", "chat", {})
     assert result == "chat_model"
     mock_initializers["special"].assert_called_once()
-    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_not_called()
+    mock_initializers["text"].assert_not_called()
 
 
 def test_chat_completion_called_when_special_returns_none(mock_initializers):
     mock_initializers["special"].return_value = None
-    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = "chat_model"
     result = init_langchain_model("chat-model", "provider", "chat", {})
     assert result == "chat_model"
     mock_initializers["special"].assert_called_once()
-    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_not_called()
+    mock_initializers["text"].assert_not_called()
 
 
 def test_community_chat_called_when_previous_fail(mock_initializers):
     mock_initializers["special"].return_value = None
-    mock_initializers["text"].return_value = None
     mock_initializers["chat"].return_value = None
     mock_initializers["community"].return_value = "community_model"
     result = init_langchain_model("community-chat", "provider", "chat", {})
     assert result == "community_model"
     mock_initializers["special"].assert_called_once()
-    mock_initializers["text"].assert_called_once()
     mock_initializers["chat"].assert_called_once()
     mock_initializers["community"].assert_called_once()
+    mock_initializers["text"].assert_not_called()
 
 
 def test_text_completion_called_when_previous_fail(mock_initializers):
@@ -204,10 +198,11 @@ def test_text_completion_called_when_previous_fail(mock_initializers):
     mock_initializers["text"].assert_called_once()
 
 
-def test_text_completion_supports_chat_mode(mock_initializers):
+def test_text_mode_only_calls_text_initializers(mock_initializers):
+    """Test that text mode only tries initializers that support text mode."""
     mock_initializers["special"].return_value = None
     mock_initializers["text"].return_value = "text_model"
-    result = init_langchain_model("text-model", "provider", "chat", {})
+    result = init_langchain_model("text-model", "provider", "text", {})
     assert result == "text_model"
     mock_initializers["special"].assert_called_once()
     mock_initializers["text"].assert_called_once()