diff --git a/tests/test_server_llm_only.py b/tests/test_server_llm_only.py new file mode 100644 index 000000000..1e799f73e --- /dev/null +++ b/tests/test_server_llm_only.py @@ -0,0 +1,308 @@ +# Copyright © 2023-2024 ValidMind Inc. All rights reserved. +# See the LICENSE file in the root of this repository for details. +# SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial + +import os +import unittest +from unittest.mock import patch + +import validmind.api_client as api_client +from validmind.ai.utils import get_client_and_model, is_configured +from validmind.tests.prompt_validation.ai_powered_test import call_model + + +class MockResponse: + def __init__(self, status, text=None, json=None): + self.status = status + self.status_code = status + self.text = text + self._json = json or {} + + def json(self): + return self._json + + +class TestServerLLMOnly(unittest.TestCase): + """Tests for server-only LLM mode functionality.""" + + def setUp(self): + """Set up test environment variables.""" + # Store original values + self.original_api_key = os.environ.get("VM_API_KEY") + self.original_api_secret = os.environ.get("VM_API_SECRET") + self.original_api_host = os.environ.get("VM_API_HOST") + self.original_api_model = os.environ.get("VM_API_MODEL") + self.original_openai_key = os.environ.get("OPENAI_API_KEY") + self.original_azure_key = os.environ.get("AZURE_OPENAI_KEY") + self.original_server_llm_only = os.environ.get("VALIDMIND_USE_SERVER_LLM_ONLY") + + # Set required environment variables for tests + os.environ["VM_API_KEY"] = "test_api_key" + os.environ["VM_API_SECRET"] = "test_api_secret" + os.environ["VM_API_HOST"] = "https://test.validmind.ai/api/v1/tracking" + os.environ["VM_API_MODEL"] = "test_model_id" + + # Clear OpenAI-related env vars to test server-only mode + if "OPENAI_API_KEY" in os.environ: + del os.environ["OPENAI_API_KEY"] + if "AZURE_OPENAI_KEY" in os.environ: + del os.environ["AZURE_OPENAI_KEY"] + if "VALIDMIND_USE_SERVER_LLM_ONLY" in os.environ: + del os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] + + def _restore_env_var(self, var_name, original_value): + """Helper to restore or clear an environment variable.""" + if original_value: + os.environ[var_name] = original_value + elif var_name in os.environ: + del os.environ[var_name] + + def _reset_ai_utils_globals(self): + """Helper to reset global state in ai.utils module.""" + import validmind.ai.utils as ai_utils + + # Reset module-level globals using name mangling pattern + # For __client in module utils, Python mangles it to _utils__client + globals_to_reset = [ + "_utils__client", + "_utils__model", + "_utils__judge_llm", + "_utils__judge_embeddings", + "_utils__ack", + ] + for global_name in globals_to_reset: + if hasattr(ai_utils, global_name): + setattr(ai_utils, global_name, None) + + def tearDown(self): + """Restore original environment variables.""" + # Restore or clear environment variables + self._restore_env_var("VM_API_KEY", self.original_api_key) + self._restore_env_var("VM_API_SECRET", self.original_api_secret) + self._restore_env_var("VM_API_HOST", self.original_api_host) + self._restore_env_var("VM_API_MODEL", self.original_api_model) + self._restore_env_var("OPENAI_API_KEY", self.original_openai_key) + self._restore_env_var("AZURE_OPENAI_KEY", self.original_azure_key) + self._restore_env_var("VALIDMIND_USE_SERVER_LLM_ONLY", self.original_server_llm_only) + + # Reset global state in ai.utils + self._reset_ai_utils_globals() + + @patch("requests.get") + def test_init_with_use_server_llm_only_parameter(self, mock_get): + """Test that use_server_llm_only parameter sets the environment variable.""" + mock_get.return_value = MockResponse( + 200, + json={ + "model": {"name": "test_model", "cuid": "test_model_id"}, + "feature_flags": {}, + "document_type": "model_documentation", + }, + ) + + api_client.init(use_server_llm_only=True) + + # Verify environment variable is set + self.assertEqual(os.environ.get("VALIDMIND_USE_SERVER_LLM_ONLY"), "True") + + # Verify ping was called + mock_get.assert_called_once() + + @patch("requests.get") + def test_init_with_use_server_llm_only_false(self, mock_get): + """Test that use_server_llm_only=False clears the environment variable.""" + # First set it to True + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = "True" + + mock_get.return_value = MockResponse( + 200, + json={ + "model": {"name": "test_model", "cuid": "test_model_id"}, + "feature_flags": {}, + "document_type": "model_documentation", + }, + ) + + api_client.init(use_server_llm_only=False) + + # Verify environment variable is set to False + self.assertEqual(os.environ.get("VALIDMIND_USE_SERVER_LLM_ONLY"), "False") + + @patch("requests.get") + def test_init_with_use_server_llm_only_none(self, mock_get): + """Test that use_server_llm_only=None doesn't change existing env var.""" + # Set it beforehand + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = "1" + + mock_get.return_value = MockResponse( + 200, + json={ + "model": {"name": "test_model", "cuid": "test_model_id"}, + "feature_flags": {}, + "document_type": "model_documentation", + }, + ) + + api_client.init(use_server_llm_only=None) + + # Verify environment variable is unchanged + self.assertEqual(os.environ.get("VALIDMIND_USE_SERVER_LLM_ONLY"), "1") + + def test_get_client_and_model_raises_error_when_server_only_enabled(self): + """Test that get_client_and_model raises ValueError when server-only mode is enabled.""" + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = "1" + + with self.assertRaises(ValueError) as context: + get_client_and_model() + + error_message = str(context.exception) + self.assertIn("Local LLM calls are disabled", error_message) + self.assertIn("server-side calls", error_message) + + def test_get_client_and_model_works_when_server_only_disabled(self): + """Test that get_client_and_model works normally when server-only mode is disabled.""" + # Set OpenAI API key to allow local calls + os.environ["OPENAI_API_KEY"] = "test_openai_key" + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = "0" + + # Should not raise an error (though it might fail on actual API call, that's OK) + # We're just testing that the server-only check doesn't block it + try: + client, model = get_client_and_model() + # If we get here, the server-only check passed + self.assertIsNotNone(client) + self.assertIsNotNone(model) + except Exception as e: + # Other errors (like API connection) are OK, we just want to ensure + # the server-only check doesn't block it + self.assertNotIn("Local LLM calls are disabled", str(e)) + + def test_is_configured_returns_true_when_server_only_enabled_and_api_configured(self): + """Test that is_configured returns True when server-only mode is enabled and API is configured.""" + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = "1" + + # Mock the API client to have credentials + with patch.object(api_client, "_api_key", "test_key"), patch.object( + api_client, "_api_secret", "test_secret" + ): + result = is_configured() + self.assertTrue(result) + + def test_is_configured_returns_false_when_server_only_enabled_but_api_not_configured(self): + """Test that is_configured returns False when server-only mode is enabled but API is not configured.""" + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = "1" + + # Mock the API client to NOT have credentials + with patch.object(api_client, "_api_key", None), patch.object( + api_client, "_api_secret", None + ): + result = is_configured() + self.assertFalse(result) + + def test_is_configured_works_normally_when_server_only_disabled(self): + """Test that is_configured works normally when server-only mode is disabled.""" + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = "0" + os.environ["OPENAI_API_KEY"] = "test_openai_key" + + # Should attempt to check OpenAI (may fail, but that's OK) + # We're just testing that server-only check doesn't interfere + try: + result = is_configured() + # Result can be True or False depending on actual API availability + self.assertIsInstance(result, bool) + except Exception: + # Other exceptions are OK, we just want to ensure server-only check doesn't block + pass + + def test_call_model_raises_error_when_server_only_enabled(self): + """Test that call_model raises ValueError when server-only mode is enabled.""" + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = "1" + + with self.assertRaises(ValueError) as context: + call_model( + system_prompt="Test system prompt", + user_prompt="Test user prompt", + ) + + error_message = str(context.exception) + self.assertIn("Local LLM calls are disabled", error_message) + self.assertIn("VALIDMIND_USE_SERVER_LLM_ONLY", error_message) + + def test_environment_variable_case_insensitive(self): + """Test that environment variable values are case-insensitive.""" + # Test various case combinations + test_cases = ["1", "True", "TRUE", "true", "0", "False", "FALSE", "false"] + + for value in test_cases: + with self.subTest(value=value): + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = value + + # Reset the ack state + import validmind.ai.utils as ai_utils + + if hasattr(ai_utils, "_utils__ack"): + setattr(ai_utils, "_utils__ack", None) + + # Check if it's treated as enabled (1, True, true, TRUE) or disabled (0, False, false, FALSE) + is_enabled = value.lower() in ["1", "true"] + + if is_enabled: + with self.assertRaises(ValueError) as context: + get_client_and_model() + self.assertIn("Local LLM calls are disabled", str(context.exception)) + else: + # When disabled, it should try to use local OpenAI (may fail, but that's OK) + try: + get_client_and_model() + except ValueError as e: + # Should not be the server-only error + self.assertNotIn("Local LLM calls are disabled", str(e)) + + @patch("requests.get") + def test_generate_descriptions_still_works_with_server_only(self, mock_get): + """Test that test result descriptions still work when server-only mode is enabled.""" + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = "1" + + mock_get.return_value = MockResponse( + 200, + json={ + "model": {"name": "test_model", "cuid": "test_model_id"}, + "feature_flags": {}, + "document_type": "model_documentation", + }, + ) + + api_client.init() + + # Test result descriptions use server-side API, so they should work + # We can't fully test this without mocking the actual API call, + # but we can verify that get_client_and_model is not called for descriptions + # (descriptions use generate_test_result_description which calls the API directly) + + # The key point is that get_client_and_model should raise an error + # but generate_test_result_description should still work + with self.assertRaises(ValueError): + get_client_and_model() + + # But we should be able to import and use the description generation + # (it uses the API client, not local OpenAI) + from validmind.ai.test_descriptions import get_result_description + + # This should work because it uses server-side API + # We can't fully test without mocking, but we can verify it doesn't + # immediately fail due to server-only mode + try: + # This will fail for other reasons (missing inputs), but not because + # of server-only mode blocking it + get_result_description( + test_id="test.test", + test_description="Test description", + should_generate=False, # Don't actually generate to avoid API call + ) + except Exception as e: + # Should not be the server-only error + self.assertNotIn("Local LLM calls are disabled", str(e)) + + +if __name__ == "__main__": + unittest.main() diff --git a/validmind/ai/utils.py b/validmind/ai/utils.py index 8f3405baf..ce310eb82 100644 --- a/validmind/ai/utils.py +++ b/validmind/ai/utils.py @@ -49,9 +49,26 @@ def get_client_and_model(): On first call, it will look in the environment for the API key endpoint, model etc. and store them in a global variable to avoid loading them up again. + + If `VALIDMIND_USE_SERVER_LLM_ONLY` is set to True, this will raise an error + indicating that local LLM calls are disabled and server-side calls should be used instead. """ global __client, __model + # Check if server-only mode is enabled + use_server_llm_only = os.getenv("VALIDMIND_USE_SERVER_LLM_ONLY", "0").lower() in [ + "1", + "true", + ] + + if use_server_llm_only: + raise ValueError( + "Local LLM calls are disabled. All LLM requests should be routed through " + "the ValidMind server. Test result descriptions already use server-side calls. " + "For prompt validation tests that require judge LLM, please ensure ValidMind AI " + "is enabled for your account and contact support if you need server-side judge LLM support." + ) + if __client and __model: return __client, __model @@ -175,11 +192,40 @@ def set_judge_config(judge_llm, judge_embeddings): def is_configured(): + """Check if LLM is configured for use. + + If server-only mode is enabled, this will check if the ValidMind API + connection is available instead of checking local OpenAI configuration. + """ global __ack if __ack: return True + # Check if server-only mode is enabled + use_server_llm_only = os.getenv("VALIDMIND_USE_SERVER_LLM_ONLY", "0").lower() in [ + "1", + "true", + ] + + if use_server_llm_only: + # In server-only mode, check if ValidMind API is connected + # by checking if we have API credentials + from ..api_client import _api_key, _api_secret + + if _api_key and _api_secret: + __ack = True + logger.debug( + "Server-only LLM mode enabled - using ValidMind API for LLM calls" + ) + return True + else: + logger.debug( + "Server-only LLM mode enabled but ValidMind API not configured" + ) + __ack = False + return False + try: client, model = get_client_and_model() # send an empty message with max_tokens=1 to "ping" the API diff --git a/validmind/api_client.py b/validmind/api_client.py index c901be491..cdbc95e9b 100644 --- a/validmind/api_client.py +++ b/validmind/api_client.py @@ -195,6 +195,7 @@ def init( model: Optional[str] = None, monitoring: bool = False, generate_descriptions: Optional[bool] = None, + use_server_llm_only: Optional[bool] = None, ): """ Initializes the API client instances and calls the /ping endpoint to ensure @@ -211,6 +212,9 @@ def init( api_host (str, optional): The API host. Defaults to None. monitoring (bool): The ongoing monitoring flag. Defaults to False. generate_descriptions (bool): Whether to use GenAI to generate test result descriptions. Defaults to True. + use_server_llm_only (bool): If True, disables local LLM calls and routes all LLM requests through + the ValidMind server. This is useful when OpenAI access is blocked locally. Defaults to None, + which respects the `VALIDMIND_USE_SERVER_LLM_ONLY` environment variable. Raises: ValueError: If the API key and secret are not provided """ @@ -239,6 +243,9 @@ def init( if generate_descriptions is not None: os.environ["VALIDMIND_LLM_DESCRIPTIONS_ENABLED"] = str(generate_descriptions) + if use_server_llm_only is not None: + os.environ["VALIDMIND_USE_SERVER_LLM_ONLY"] = str(use_server_llm_only) + reload() diff --git a/validmind/tests/prompt_validation/ai_powered_test.py b/validmind/tests/prompt_validation/ai_powered_test.py index 03ce32cfa..5fda75834 100644 --- a/validmind/tests/prompt_validation/ai_powered_test.py +++ b/validmind/tests/prompt_validation/ai_powered_test.py @@ -29,6 +29,22 @@ def call_model( judge_embeddings=None, ): """Call LLM with the given prompts and return the response""" + import os + + use_server_llm_only = os.getenv("VALIDMIND_USE_SERVER_LLM_ONLY", "0").lower() in [ + "1", + "true", + ] + + if use_server_llm_only: + raise ValueError( + "Local LLM calls are disabled (VALIDMIND_USE_SERVER_LLM_ONLY is enabled). " + "Prompt validation tests that require judge LLM currently need local OpenAI access. " + "Please either:\n" + "1. Disable server-only mode if you have OpenAI access, or\n" + "2. Contact ValidMind support to enable server-side judge LLM support for your account." + ) + if not is_configured(): raise ValueError( "LLM is not configured. Please set an `OPENAI_API_KEY` environment variable "