Fix tests

Author: JashG

nemoguardrails/actions/llm/utils.py

@@ -47,6 +47,7 @@
     "endpoint",
     "endpoint_url",
     "openai_api_base",
+    "server_url",
 ]
 
 

nemoguardrails/rails/llm/config.py

@@ -1526,7 +1526,7 @@ def check_model_exists_for_output_rails(cls, values):
                     else "none"
                 )
                 raise InvalidRailsConfigurationError(
-                    "Output flow {flow_id} references model type '{flow_model}' that is not defined in the configuration. Detected model types: {available_types}."
+                    f"Output flow '{flow_id}' references model type '{flow_model}' that is not defined in the configuration. Detected model types: {available_types}."
                 )
         return values
 

tests/test_actions_llm_utils.py

@@ -13,9 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import cast
 from unittest.mock import AsyncMock
 
 import pytest
+from langchain_core.language_models import BaseLanguageModel
 
 from nemoguardrails.actions.llm.utils import (
     _extract_and_remove_think_tags,
@@ -55,6 +57,24 @@ 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)
@@ -315,14 +335,13 @@ def test_extract_and_remove_think_tags_wrong_order():
 @pytest.mark.asyncio
 async def test_llm_call_exception_enrichment_with_model_and_endpoint():
     """Test that LLM invocation errors include model and endpoint context."""
-    mock_llm = AsyncMock()
-    mock_llm.__module__ = "langchain_openai.chat_models"
+    mock_llm = MockOpenAILLM()
     mock_llm.model_name = "gpt-4"
     mock_llm.base_url = "https://api.openai.com/v1"
     mock_llm.ainvoke = AsyncMock(side_effect=ConnectionError("Connection refused"))
 
     with pytest.raises(LLMCallException) as exc_info:
-        await llm_call(mock_llm, "test prompt")
+        await llm_call(cast(BaseLanguageModel, mock_llm), "test prompt")
 
     exc_str = str(exc_info.value)
     assert "gpt-4" in exc_str
@@ -351,14 +370,13 @@ async def test_llm_call_exception_without_endpoint():
 @pytest.mark.asyncio
 async def test_llm_call_exception_extracts_azure_endpoint():
     """Test that Azure-style endpoint URLs are extracted."""
-    mock_llm = AsyncMock()
-    mock_llm.__module__ = "langchain_openai.chat_models"
+    mock_llm = MockAzureLLM()
     mock_llm.model_name = "gpt-4"
     mock_llm.azure_endpoint = "https://example.openai.azure.com"
     mock_llm.ainvoke = AsyncMock(side_effect=Exception("Azure error"))
 
     with pytest.raises(LLMCallException) as exc_info:
-        await llm_call(mock_llm, "test prompt")
+        await llm_call(cast(BaseLanguageModel, mock_llm), "test prompt")
 
     exc_str = str(exc_info.value)
     assert "https://example.openai.azure.com" in exc_str
@@ -369,14 +387,13 @@ async def test_llm_call_exception_extracts_azure_endpoint():
 @pytest.mark.asyncio
 async def test_llm_call_exception_extracts_server_url():
     """Test that TRT-style server_url is extracted."""
-    mock_llm = AsyncMock()
-    mock_llm.__module__ = "nemoguardrails.llm.providers.trtllm.llm"
+    mock_llm = MockTRTLLM()
     mock_llm.model_name = "llama-2-70b"
     mock_llm.server_url = "https://triton.example.com:8000"
     mock_llm.ainvoke = AsyncMock(side_effect=Exception("Triton server error"))
 
     with pytest.raises(LLMCallException) as exc_info:
-        await llm_call(mock_llm, "test prompt")
+        await llm_call(cast(BaseLanguageModel, mock_llm), "test prompt")
 
     exc_str = str(exc_info.value)
     assert "https://triton.example.com:8000" in exc_str
@@ -387,18 +404,12 @@ async def test_llm_call_exception_extracts_server_url():
 @pytest.mark.asyncio
 async def test_llm_call_exception_extracts_nested_client_base_url():
     """Test that nested client.base_url is extracted."""
-
-    class MockClient:
-        base_url = "https://custom.endpoint.com/v1"
-
-    mock_llm = AsyncMock()
-    mock_llm.__module__ = "langchain_openai.chat_models"
+    mock_llm = MockLLMWithClient()
     mock_llm.model_name = "gpt-4-turbo"
-    mock_llm.client = MockClient()
     mock_llm.ainvoke = AsyncMock(side_effect=Exception("Client error"))
 
     with pytest.raises(LLMCallException) as exc_info:
-        await llm_call(mock_llm, "test prompt")
+        await llm_call(cast(BaseLanguageModel, mock_llm), "test prompt")
 
     exc_str = str(exc_info.value)
     assert "https://custom.endpoint.com/v1" in exc_str

tests/test_config_validation.py

@@ -75,7 +75,7 @@ def test_self_check_input_prompt_exception():
         )
         LLMRails(config=config)
 
-    assert "You must provide a `self_check_input` prompt" in str(exc_info.value)
+    assert "Missing a `self_check_input` prompt template" in str(exc_info.value)
 
 
 def test_self_check_output_prompt_exception():
@@ -90,7 +90,7 @@ def test_self_check_output_prompt_exception():
         )
         LLMRails(config=config)
 
-    assert "You must provide a `self_check_output` prompt" in str(exc_info.value)
+    assert "Missing a `self_check_output` prompt template" in str(exc_info.value)
 
 
 def test_passthrough_and_single_call_incompatibility():