mito-ai 0.1.57__py3-none-any.whl → 0.1.58__py3-none-any.whl

mito_ai/tests/providers/utils.py

@@ -38,7 +38,7 @@ def mock_openai_client() -> Any:
     mock_client.key_type = "user"
     mock_client.request_completions = AsyncMock(return_value="Test completion")
     mock_client.stream_completions = AsyncMock(return_value="Test completion")
-    return patch("mito_ai.completions.providers.OpenAIClient", return_value=mock_client)
+    return patch("mito_ai.provider_manager.OpenAIClient", return_value=mock_client)
 
 
 def mock_gemini_client() -> Any:
@@ -52,7 +52,7 @@ def mock_gemini_client() -> Any:
     mock_client.key_type = "gemini"
     mock_client.request_completions = AsyncMock(return_value="Test completion")
     mock_client.stream_completions = AsyncMock(return_value="Test completion")
-    return patch("mito_ai.completions.providers.GeminiClient", return_value=mock_client)
+    return patch("mito_ai.provider_manager.GeminiClient", return_value=mock_client)
 
 
 def mock_azure_openai_client() -> Any:
@@ -66,7 +66,7 @@ def mock_azure_openai_client() -> Any:
     mock_client.key_type = "azure"
     mock_client.request_completions = AsyncMock(return_value="Test completion")
     mock_client.stream_completions = AsyncMock(return_value="Test completion")
-    return patch("mito_ai.completions.providers.OpenAIClient", return_value=mock_client)
+    return patch("mito_ai.provider_manager.OpenAIClient", return_value=mock_client)
 
 
 
@@ -82,4 +82,4 @@ def mock_claude_client() -> Any:
     mock_client.request_completions = AsyncMock(return_value="Test completion")
     mock_client.stream_completions = AsyncMock(return_value="Test completion")
     mock_client.stream_response = AsyncMock(return_value="Test completion")
-    return patch("mito_ai.completions.providers.AnthropicClient", return_value=mock_client)
+    return patch("mito_ai.provider_manager.AnthropicClient", return_value=mock_client)

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -2,94 +2,38 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 from typing import List
-from anthropic.types import MessageParam
 import pytest
 import os
 from unittest.mock import patch, AsyncMock, MagicMock
+from traitlets.config import Config
 from mito_ai.streamlit_conversion.streamlit_agent_handler import (
-    get_response_from_agent,
     generate_new_streamlit_code,
     correct_error_in_generation,
     streamlit_handler
 )
 from mito_ai.path_utils import AbsoluteNotebookPath, AppFileName, get_absolute_app_path, get_absolute_notebook_dir_path, get_absolute_notebook_path
+from mito_ai.provider_manager import ProviderManager
 
 # Add this line to enable async support
 pytest_plugins = ('pytest_asyncio',)
 
 
-class TestGetResponseFromAgent:
-    """Test cases for get_response_from_agent function"""
-
-    @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
-    async def test_get_response_from_agent_success(self, mock_stream):
-        """Test get_response_from_agent with successful response"""
-        # Mock the async generator
-        async def mock_async_gen():
-            yield "Here's your code:\n```python\nimport streamlit\nst.title('Test')\n```"
-
-        mock_stream.return_value = mock_async_gen()
-        
-        messages: List[MessageParam] = [{"role": "user", "content": [{"type": "text", "text": "test"}]}]
-        response = await get_response_from_agent(messages)
-        
-        assert response is not None
-        assert len(response) > 0
-        assert "import streamlit" in response
-
-    @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
-    @pytest.mark.parametrize("mock_items,expected_result", [
-        (["Hello", " World", "!"], "Hello World!"),
-        ([], ""),
-        (["Here's your code: import streamlit"], "Here's your code: import streamlit")
-    ])
-    async def test_get_response_from_agent_parametrized(self, mock_stream, mock_items, expected_result):
-        """Test response from agent with different scenarios"""
-        # Mock the async generator
-        async def mock_async_gen():
-            for item in mock_items:
-                yield item
-
-        mock_stream.return_value = mock_async_gen()
-        
-        messages: List[MessageParam] = [{"role": "user", "content": [{"type": "text", "text": "test"}]}]
-        result = await get_response_from_agent(messages)
-        
-        assert result == expected_result
-        mock_stream.assert_called_once()
-        
-
-    @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
-    async def test_get_response_from_agent_exception(self, mock_stream):
-        """Test exception handling in get_response_from_agent"""
-        mock_stream.side_effect = Exception("API Error")
-        
-        messages: List[MessageParam] = [{"role": "user", "content": [{"type": "text", "text": "test"}]}]
-        
-        with pytest.raises(Exception, match="API Error"):
-            await get_response_from_agent(messages)
-
-
 class TestGenerateStreamlitCode:
     """Test cases for generate_new_streamlit_code function"""
 
     @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
-    async def test_generate_new_streamlit_code_success(self, mock_stream):
+    async def test_generate_new_streamlit_code_success(self):
         """Test successful streamlit code generation"""
         mock_response = "Here's your code:\n```python\nimport streamlit\nst.title('Hello')\n```"
 
-        async def mock_async_gen():
-            for item in [mock_response]:
-                yield item
-
-        mock_stream.return_value = mock_async_gen()
+        provider_config = Config()
+        provider_config.ProviderManager = Config()
+        provider_config.OpenAIClient = Config()
+        provider = ProviderManager(config=provider_config)
         
-        notebook_data: List[dict] = [{"cells": []}]
-        result = await generate_new_streamlit_code(notebook_data, '')
+        with patch.object(provider, 'request_completions', new_callable=AsyncMock, return_value=mock_response):
+            notebook_data: List[dict] = [{"cells": []}]
+            result = await generate_new_streamlit_code(notebook_data, '', provider)
         
         expected_code = "import streamlit\nst.title('Hello')\n"
         assert result == expected_code
@@ -99,8 +43,7 @@ class TestCorrectErrorInGeneration:
     """Test cases for correct_error_in_generation function"""
 
     @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
-    async def test_correct_error_in_generation_success(self, mock_stream):
+    async def test_correct_error_in_generation_success(self):
         """Test successful error correction"""
         mock_response = """```search_replace
 >>>>>>> SEARCH
@@ -109,25 +52,28 @@ st.title('Test')
 st.title('Fixed')
 <<<<<<< REPLACE
 ```"""
-        async def mock_async_gen():
-            for item in [mock_response]:
-                yield item
-
-        mock_stream.return_value = mock_async_gen()
+        provider_config = Config()
+        provider_config.ProviderManager = Config()
+        provider_config.OpenAIClient = Config()
+        provider = ProviderManager(config=provider_config)
 
-        result = await correct_error_in_generation("ImportError: No module named 'pandas'", "import streamlit\nst.title('Test')\n")
+        with patch.object(provider, 'request_completions', new_callable=AsyncMock, return_value=mock_response):
+            result = await correct_error_in_generation("ImportError: No module named 'pandas'", "import streamlit\nst.title('Test')\n", provider)
         
         expected_code = "import streamlit\nst.title('Fixed')\n"
         assert result == expected_code
 
     @pytest.mark.asyncio
-    @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
-    async def test_correct_error_in_generation_exception(self, mock_stream):
+    async def test_correct_error_in_generation_exception(self):
         """Test exception handling in error correction"""
-        mock_stream.side_effect = Exception("API Error")
+        provider_config = Config()
+        provider_config.ProviderManager = Config()
+        provider_config.OpenAIClient = Config()
+        provider = ProviderManager(config=provider_config)
         
-        with pytest.raises(Exception, match="API Error"):
-            await correct_error_in_generation("Some error", "import streamlit\nst.title('Test')")
+        with patch.object(provider, 'request_completions', new_callable=AsyncMock, side_effect=Exception("API Error")):
+            with pytest.raises(Exception, match="API Error"):
+                await correct_error_in_generation("Some error", "import streamlit\nst.title('Test')", provider)
 
 
 class TestStreamlitHandler:
@@ -158,11 +104,17 @@ class TestStreamlitHandler:
         # Construct the expected app path using the same method as the production code
         app_directory = get_absolute_notebook_dir_path(notebook_path)
         expected_app_path = get_absolute_app_path(app_directory, app_file_name)
-        await streamlit_handler(True, notebook_path, app_file_name)
+        
+        provider_config = Config()
+        provider_config.ProviderManager = Config()
+        provider_config.OpenAIClient = Config()
+        provider = ProviderManager(config=provider_config)
+        
+        await streamlit_handler(True, notebook_path, app_file_name, '', provider)
         
         # Verify calls
         mock_parse.assert_called_once_with(notebook_path)
-        mock_generate_code.assert_called_once_with(mock_notebook_data, '')
+        mock_generate_code.assert_called_once_with(mock_notebook_data, '', provider)
         mock_validator.assert_called_once_with("import streamlit\nst.title('Test')", notebook_path)
         mock_create_file.assert_called_once_with(expected_app_path, "import streamlit\nst.title('Test')")
 
@@ -185,9 +137,14 @@ class TestStreamlitHandler:
         # Mock validation (always errors) - validate_app returns List[str]
         mock_validator.return_value = ["Persistent error"]
     
+        provider_config = Config()
+        provider_config.ProviderManager = Config()
+        provider_config.OpenAIClient = Config()
+        provider = ProviderManager(config=provider_config)
+    
         # Now it should raise an exception instead of returning a tuple
         with pytest.raises(Exception):
-            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '')
+            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '', provider)
         
         # Verify that error correction was called 5 times (once per error, 5 retries)
         # Each retry processes 1 error, so 5 retries = 5 calls
@@ -213,9 +170,14 @@ class TestStreamlitHandler:
         # Mock file creation failure - now it should raise an exception
         mock_create_file.side_effect = Exception("Permission denied")
         
+        provider_config = Config()
+        provider_config.ProviderManager = Config()
+        provider_config.OpenAIClient = Config()
+        provider = ProviderManager(config=provider_config)
+        
         # Now it should raise an exception instead of returning a tuple
         with pytest.raises(Exception):
-            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '')
+            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '', provider)
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
@@ -224,8 +186,13 @@ class TestStreamlitHandler:
         
         mock_parse.side_effect = FileNotFoundError("Notebook not found")
         
+        provider_config = Config()
+        provider_config.ProviderManager = Config()
+        provider_config.OpenAIClient = Config()
+        provider = ProviderManager(config=provider_config)
+        
         with pytest.raises(FileNotFoundError, match="Notebook not found"):
-            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '')
+            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '', provider)
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
@@ -239,8 +206,13 @@ class TestStreamlitHandler:
         # Mock code generation failure
         mock_generate_code.side_effect = Exception("Generation failed")
         
+        provider_config = Config()
+        provider_config.ProviderManager = Config()
+        provider_config.OpenAIClient = Config()
+        provider = ProviderManager(config=provider_config)
+        
         with pytest.raises(Exception, match="Generation failed"):
-            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '')
+            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '', provider)
 
 
 

mito_ai/tests/streamlit_preview/test_streamlit_preview_handler.py

@@ -63,12 +63,15 @@ class TestStreamlitPreviewHandler:
             mock_request.connection = MagicMock()
             mock_request.connection.context = MagicMock()
             
+            # Create a mock ProviderManager for the llm parameter
+            mock_llm = MagicMock()
+
             # Create handler instance  
             handler = StreamlitPreviewHandler(
                 application=mock_application,
                 request=mock_request,
+                llm=mock_llm
             )
-            handler.initialize()
             
             # Mock authentication - set current_user to bypass @tornado.web.authenticated
             handler.current_user = "test_user"  # type: ignore

mito_ai/tests/test_enterprise_mode.py

@@ -0,0 +1,162 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+from unittest.mock import patch, MagicMock
+from traitlets.config import Config
+from mito_ai.utils.telemetry_utils import telemetry_turned_on, identify, log
+from mito_ai.utils.model_utils import get_available_models
+from mito_ai.provider_manager import ProviderManager
+from mito_ai.completions.models import MessageType
+from openai.types.chat import ChatCompletionMessageParam
+
+
+@pytest.fixture
+def provider_config() -> Config:
+    """Create a proper Config object for the ProviderManager."""
+    config = Config()
+    config.ProviderManager = Config()
+    config.OpenAIClient = Config()
+    return config
+
+
+class TestEnterpriseModeDetection:
+    """Tests for enterprise mode detection."""
+    
+    @patch('mito_ai.utils.version_utils.is_enterprise')
+    def test_telemetry_disabled_when_enterprise(self, mock_is_enterprise):
+        """Test that telemetry is disabled when enterprise mode is enabled."""
+        mock_is_enterprise.return_value = True
+        
+        result = telemetry_turned_on()
+        
+        assert result is False
+    
+    @patch('mito_ai.utils.telemetry_utils.is_enterprise')
+    def test_telemetry_enabled_when_not_enterprise(self, mock_is_enterprise):
+        """Test that telemetry can be enabled when enterprise mode is not enabled."""
+        mock_is_enterprise.return_value = False
+        
+        # Mock other conditions that might disable telemetry
+        with patch('mito_ai.utils.telemetry_utils.MITOSHEET_HELPER_PRIVATE', False), \
+             patch('mito_ai.utils.telemetry_utils.is_pro', return_value=False), \
+             patch('mito_ai.utils.telemetry_utils.get_user_field', return_value=True):
+            result = telemetry_turned_on()
+            # Result depends on other conditions, but enterprise check should pass
+            # We just verify enterprise check doesn't block it
+            mock_is_enterprise.assert_called_once()
+
+
+class TestTelemetryDisabling:
+    """Tests for telemetry disabling in enterprise mode."""
+    
+    @patch('mito_ai.utils.version_utils.is_enterprise')
+    @patch('mito_ai.utils.telemetry_utils.analytics')
+    def test_identify_skips_when_enterprise(self, mock_analytics, mock_is_enterprise):
+        """Test that identify() skips analytics calls when enterprise mode is enabled."""
+        mock_is_enterprise.return_value = True
+        
+        identify()
+        
+        # Should not call analytics.identify
+        mock_analytics.identify.assert_not_called()
+    
+    @patch('mito_ai.utils.version_utils.is_enterprise')
+    @patch('mito_ai.utils.telemetry_utils.analytics')
+    def test_log_skips_when_enterprise(self, mock_analytics, mock_is_enterprise):
+        """Test that log() skips analytics calls when enterprise mode is enabled."""
+        mock_is_enterprise.return_value = True
+        
+        log("test_event", {"param": "value"})
+        
+        # Should not call analytics.track
+        mock_analytics.track.assert_not_called()
+
+
+class TestModelValidation:
+    """Tests for model validation in enterprise mode."""
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    def test_provider_manager_validates_model(self, mock_constants, mock_is_enterprise, provider_config: Config):
+        """Test that ProviderManager validates models against available models."""
+        mock_is_enterprise.return_value = True
+        mock_constants.LITELLM_BASE_URL = "https://litellm-server.com"
+        mock_constants.LITELLM_MODELS = ["openai/gpt-4o", "openai/gpt-4o-mini"]
+        
+        provider_manager = ProviderManager(config=provider_config)
+        provider_manager.set_selected_model("openai/gpt-4o")
+        
+        # Should not raise an error for valid model
+        available_models = get_available_models()
+        assert "openai/gpt-4o" in available_models
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    @pytest.mark.asyncio
+    async def test_provider_manager_rejects_invalid_model(self, mock_constants, mock_is_enterprise, provider_config: Config):
+        """Test that ProviderManager rejects invalid models."""
+        mock_is_enterprise.return_value = True
+        mock_constants.LITELLM_BASE_URL = "https://litellm-server.com"
+        mock_constants.LITELLM_MODELS = ["openai/gpt-4o"]
+        mock_constants.LITELLM_API_KEY = "test-key"
+        
+        provider_manager = ProviderManager(config=provider_config)
+        provider_manager.set_selected_model("invalid-model")
+        
+        messages: list[ChatCompletionMessageParam] = [{"role": "user", "content": "test"}]
+        
+        # Should raise ValueError for invalid model
+        with pytest.raises(ValueError, match="is not in the allowed model list"):
+            await provider_manager.request_completions(
+                message_type=MessageType.CHAT,
+                messages=messages
+            )
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    def test_available_models_endpoint_returns_litellm_models(self, mock_constants, mock_is_enterprise):
+        """Test that /available-models endpoint returns LiteLLM models when configured."""
+        mock_is_enterprise.return_value = True
+        mock_constants.LITELLM_BASE_URL = "https://litellm-server.com"
+        mock_constants.LITELLM_MODELS = ["openai/gpt-4o", "anthropic/claude-3-5-sonnet"]
+        
+        result = get_available_models()
+        
+        assert result == ["openai/gpt-4o", "anthropic/claude-3-5-sonnet"]
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    def test_available_models_endpoint_returns_standard_models_when_not_configured(self, mock_constants, mock_is_enterprise):
+        """Test that /available-models endpoint returns standard models when LiteLLM is not configured."""
+        mock_is_enterprise.return_value = True
+        mock_constants.LITELLM_BASE_URL = None
+        mock_constants.LITELLM_MODELS = []
+        
+        result = get_available_models()
+        
+        from mito_ai.utils.model_utils import STANDARD_MODELS
+        assert result == STANDARD_MODELS
+
+
+class TestModelStorage:
+    """Tests for model storage in ProviderManager."""
+    
+    def test_provider_manager_stores_model(self, provider_config: Config):
+        """Test that ProviderManager can store and retrieve selected model."""
+        provider_manager = ProviderManager(config=provider_config)
+        
+        provider_manager.set_selected_model("gpt-4.1")
+        assert provider_manager.get_selected_model() == "gpt-4.1"
+        
+        provider_manager.set_selected_model("claude-sonnet-4-5-20250929")
+        assert provider_manager.get_selected_model() == "claude-sonnet-4-5-20250929"
+    
+    def test_provider_manager_default_model(self, provider_config: Config):
+        """Test that ProviderManager has a default model."""
+        provider_manager = ProviderManager(config=provider_config)
+        
+        # Should have default model
+        default_model = provider_manager.get_selected_model()
+        assert default_model is not None
+        assert isinstance(default_model, str)