mito-ai 0.1.37__py3-none-any.whl → 0.1.39__py3-none-any.whl

mito_ai/tests/streamlit_preview/test_streamlit_preview_manager.py

@@ -0,0 +1,302 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import time
+import tempfile
+import os
+import shutil
+import subprocess
+import threading
+import requests
+from unittest.mock import Mock, patch, MagicMock
+from typing import Any
+
+from mito_ai.streamlit_preview.manager import (
+    StreamlitPreviewManager, 
+    PreviewProcess, 
+    get_preview_manager
+)
+
+
+class TestStreamlitPreviewManager:
+    """Test cases for StreamlitPreviewManager."""
+    
+    @pytest.fixture
+    def manager(self):
+        """Create a fresh manager instance for each test."""
+        return StreamlitPreviewManager()
+    
+    @pytest.fixture
+    def sample_app_code(self):
+        """Sample streamlit app code for testing."""
+        return """
+import streamlit as st
+
+st.title("Test App")
+st.write("Hello, World!")
+"""
+    
+    def test_init(self, manager):
+        """Test manager initialization."""
+        assert manager._previews == {}
+        assert isinstance(manager._lock, type(threading.Lock()))
+        assert manager.log is not None
+    
+    def test_get_free_port(self, manager):
+        """Test getting a free port."""
+        port = manager.get_free_port()
+        assert isinstance(port, int)
+        assert port > 0
+        assert port < 65536
+        
+        # Test that we get different ports
+        port2 = manager.get_free_port()
+        assert port != port2
+    
+    @pytest.mark.parametrize("app_code,preview_id,expected_success", [
+        ("import streamlit as st\nst.write('Hello')", "test_preview", True),
+        ("", "empty_preview", True),
+        ("import streamlit as st\n" * 1000 + "st.write('Large app')", "large_preview", True),
+    ])
+    def test_start_streamlit_preview_success_cases(self, manager, app_code, preview_id, expected_success):
+        """Test successful streamlit preview start with different app codes."""
+        with patch('subprocess.Popen') as mock_popen, \
+             patch('requests.get') as mock_requests_get, \
+             patch('tempfile.mkdtemp') as mock_mkdtemp:
+            
+            # Setup mocks
+            app_directory = "/tmp/test_dir"
+            mock_mkdtemp.return_value = app_directory
+            mock_proc = Mock()
+            mock_proc.terminate.return_value = None
+            mock_proc.wait.return_value = None
+            mock_popen.return_value = mock_proc
+            
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_requests_get.return_value = mock_response
+            
+            # Test
+            success, message, port = manager.start_streamlit_preview(app_directory, preview_id)
+            
+            # Assertions
+            assert success == expected_success
+            if expected_success:
+                assert "successfully" in message.lower()
+                assert isinstance(port, int)
+                assert port > 0
+                
+                # Verify subprocess was called correctly
+                mock_popen.assert_called_once()
+                call_args = mock_popen.call_args
+                assert "streamlit" in call_args[0][0]
+                assert "run" in call_args[0][0]
+                assert "--server.headless" in call_args[0][0]
+                assert "--server.address" in call_args[0][0]
+                
+                # Cleanup
+                manager.stop_preview(preview_id)
+    
+    @pytest.mark.parametrize("exception_type,expected_message", [
+        (Exception("Temp dir creation failed"), "failed to start preview"),
+        (OSError("Permission denied"), "failed to start preview"),
+        (ValueError("Invalid argument"), "failed to start preview"),
+    ])
+    def test_start_streamlit_preview_exceptions(self, manager, sample_app_code, exception_type, expected_message):
+        """Test streamlit preview start with different exceptions."""
+        with patch('tempfile.mkdtemp', side_effect=exception_type):
+            app_directory = "/tmp/test_dir"
+            success, message, port = manager.start_streamlit_preview(app_directory, "test_preview")
+            
+            assert success is False
+            assert expected_message in message.lower()
+            assert port is None
+    
+    @pytest.mark.parametrize("preview_id,expected_result", [
+        ("existing_preview", True),
+        ("non_existent", False),
+    ])
+    def test_stop_preview_scenarios(self, manager, sample_app_code, preview_id, expected_result):
+        """Test stopping previews with different scenarios."""
+        if expected_result:
+            # Start a preview first
+            with patch('subprocess.Popen') as mock_popen, \
+                 patch('requests.get') as mock_requests_get, \
+                 patch('tempfile.mkdtemp') as mock_mkdtemp, \
+                 patch('builtins.open', create=True) as mock_open, \
+                 patch('os.path.exists') as mock_exists:
+                
+                app_directory = "/tmp/test_dir"
+                mock_mkdtemp.return_value = app_directory
+                mock_proc = Mock()
+                mock_proc.terminate.return_value = None
+                mock_proc.wait.return_value = None
+                mock_popen.return_value = mock_proc
+                
+                mock_response = Mock()
+                mock_response.status_code = 200
+                mock_requests_get.return_value = mock_response
+                
+                # Mock file operations
+                mock_file = Mock()
+                mock_open.return_value.__enter__.return_value = mock_file
+                mock_exists.return_value = True
+                
+                manager.start_streamlit_preview(app_directory, preview_id)
+    
+    @pytest.mark.parametrize("process_behavior,expected_kill_called", [
+        (subprocess.TimeoutExpired("cmd", 5), True),
+        (None, False),  # Normal termination
+    ])
+    def test_stop_preview_process_behaviors(self, manager, sample_app_code, process_behavior, expected_kill_called):
+        """Test stopping preview with different process behaviors."""
+        with patch('subprocess.Popen') as mock_popen, \
+             patch('requests.get') as mock_requests_get, \
+             patch('tempfile.mkdtemp') as mock_mkdtemp, \
+             patch('builtins.open', create=True) as mock_open, \
+             patch('os.path.exists') as mock_exists:
+            
+            # Setup mocks for start
+            app_directory = "/tmp/test_dir"
+            mock_mkdtemp.return_value = app_directory
+            
+            mock_proc = Mock()
+            mock_proc.terminate.return_value = None
+            mock_proc.wait.return_value = None
+            mock_popen.return_value = mock_proc
+            
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_requests_get.return_value = mock_response
+            
+            # Mock file operations
+            mock_file = Mock()
+            mock_open.return_value.__enter__.return_value = mock_file
+            mock_exists.return_value = True
+            
+            # Start a preview
+            manager.start_streamlit_preview(app_directory, "test_preview")
+            
+            # Setup process behavior for stop
+            if process_behavior:
+                # Configure the mock to raise the exception when called with timeout
+                def wait_with_timeout(*args, **kwargs):
+                    if 'timeout' in kwargs:
+                        raise process_behavior
+                    return None
+                mock_proc.wait.side_effect = wait_with_timeout
+    
+    @pytest.mark.parametrize("preview_id,expected_found", [
+        ("existing_preview", True),
+        ("non_existent", False),
+    ])
+    def test_get_preview_scenarios(self, manager, sample_app_code, preview_id, expected_found):
+        """Test getting previews with different scenarios."""
+        if expected_found:
+            # Start a preview first
+            with patch('subprocess.Popen') as mock_popen, \
+                 patch('requests.get') as mock_requests_get, \
+                 patch('tempfile.mkdtemp') as mock_mkdtemp, \
+                 patch('builtins.open', create=True) as mock_open, \
+                 patch('os.path.exists') as mock_exists:
+                
+                mock_mkdtemp.return_value = "/tmp/test_dir"
+                mock_proc = Mock()
+                mock_proc.terminate.return_value = None
+                mock_proc.wait.return_value = None
+                mock_popen.return_value = mock_proc
+                
+                mock_response = Mock()
+                mock_response.status_code = 200
+                mock_requests_get.return_value = mock_response
+                
+                # Mock file operations
+                mock_file = Mock()
+                mock_open.return_value.__enter__.return_value = mock_file
+                mock_exists.return_value = True
+                
+                manager.start_streamlit_preview("/tmp/test_dir", preview_id)
+        
+        preview = manager.get_preview(preview_id)
+        
+        if expected_found:
+            assert preview is not None
+            assert isinstance(preview, PreviewProcess)
+            assert preview.port > 0
+            
+            # Cleanup
+            manager.stop_preview(preview_id)
+        else:
+            assert preview is None
+    
+    def test_preview_process_dataclass(self):
+        """Test PreviewProcess dataclass."""
+        proc = Mock()
+        port = 8080
+        
+        preview = PreviewProcess(
+            proc=proc,
+            port=port
+        )
+        
+        assert preview.proc == proc
+        assert preview.port == port
+    
+    def test_get_preview_manager_singleton(self):
+        """Test that get_preview_manager returns the same instance."""
+        manager1 = get_preview_manager()
+        manager2 = get_preview_manager()
+        
+        assert manager1 is manager2
+        assert isinstance(manager1, StreamlitPreviewManager)
+    
+    @pytest.mark.parametrize("num_previews", [1, 2, 3])
+    def test_concurrent_previews(self, manager, sample_app_code, num_previews):
+        """Test managing multiple concurrent previews."""
+        preview_ids = [f"preview_{i}" for i in range(num_previews)]
+        ports = []
+        
+        with patch('subprocess.Popen') as mock_popen, \
+             patch('requests.get') as mock_requests_get, \
+             patch('tempfile.mkdtemp') as mock_mkdtemp, \
+             patch('builtins.open', create=True) as mock_open, \
+             patch('os.path.exists') as mock_exists:
+            
+            # Setup mocks
+            mock_mkdtemp.return_value = "/tmp/test_dir"
+            mock_proc = Mock()
+            mock_proc.terminate.return_value = None
+            mock_proc.wait.return_value = None
+            mock_popen.return_value = mock_proc
+            
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_requests_get.return_value = mock_response
+            
+            # Mock file operations
+            mock_file = Mock()
+            mock_open.return_value.__enter__.return_value = mock_file
+            mock_exists.return_value = True
+            
+            # Start multiple previews
+            for preview_id in preview_ids:
+                success, _, port = manager.start_streamlit_preview("/tmp/test_dir", preview_id)
+                assert success is True
+                ports.append(port)
+            
+            # Assertions
+            assert len(set(ports)) == num_previews  # All ports should be different
+            
+            # Check all previews exist
+            for preview_id in preview_ids:
+                assert manager.get_preview(preview_id) is not None
+            
+            # Stop all previews
+            for preview_id in preview_ids:
+                assert manager.stop_preview(preview_id)
+            
+            # Verify they're gone
+            for preview_id in preview_ids:
+                assert manager.get_preview(preview_id) is None
+    

mito_ai/tests/utils/test_anthropic_utils.py

@@ -5,7 +5,7 @@ import pytest
 import anthropic
 from typing import List, Dict, Any, Tuple, Union, cast
 from anthropic.types import MessageParam, ToolUnionParam, ToolParam
-from mito_ai.utils.anthropic_utils import _prepare_anthropic_request_data_and_headers
+from mito_ai.utils.anthropic_utils import ANTHROPIC_TIMEOUT, _prepare_anthropic_request_data_and_headers
 from mito_ai.completions.models import MessageType
 from mito_ai.utils.schema import UJ_STATIC_USER_ID, UJ_USER_EMAIL
 from mito_ai.utils.db import get_user_field
@@ -52,7 +52,7 @@ def test_basic_request_preparation():
     )
 
     assert headers == {"Content-Type": "application/json"}
-    assert data["timeout"] == 30
+    assert data["timeout"] == ANTHROPIC_TIMEOUT
     assert data["max_retries"] == 1
     assert data["email"] == "test@example.com"
     assert data["user_id"] == "test_user_id"

mito_ai/utils/anthropic_utils.py

@@ -23,7 +23,7 @@ from mito_ai.constants import MITO_ANTHROPIC_URL
 __user_email: Optional[str] = None
 __user_id: Optional[str] = None
 
-timeout = 30
+ANTHROPIC_TIMEOUT = 60
 max_retries = 1
 
 FAST_ANTHROPIC_MODEL = "claude-3-5-haiku-latest"
@@ -63,7 +63,7 @@ def _prepare_anthropic_request_data_and_headers(
         inner_data["stream"] = stream
     # Compose the outer data dict
     data = {
-        "timeout": timeout,
+        "timeout": ANTHROPIC_TIMEOUT,
         "max_retries": max_retries,
         "email": __user_email,
         "user_id": __user_id,
@@ -90,7 +90,7 @@ async def get_anthropic_completion_from_mito_server(
         MITO_ANTHROPIC_URL, 
         headers, 
         data, 
-        timeout, 
+        ANTHROPIC_TIMEOUT, 
         max_retries, 
         message_type, 
         provider_name="Claude"
@@ -118,7 +118,7 @@ async def stream_anthropic_completion_from_mito_server(
         url=MITO_ANTHROPIC_URL,
         headers=headers,
         data=data,
-        timeout=timeout,
+        timeout=ANTHROPIC_TIMEOUT,
         max_retries=max_retries,
         message_type=message_type,
         reply_fn=actual_reply_fn,

mito_ai/utils/open_ai_utils.py

@@ -196,9 +196,5 @@ def get_open_ai_completion_function_params(
                 "strict": True
             }
         }
-    
-    # o3-mini will error if we try setting the temperature
-    if not model.startswith("o3"):
-        completion_function_params["temperature"] = 0.0
 
     return completion_function_params

mito_ai/utils/telemetry_utils.py

@@ -174,7 +174,6 @@ def log(
     If telemetry is not turned off and we are not running tests,
     we log the ai event
     """
-
     final_params: Dict[str, Any] = params or {}
     
     # Then, make sure to add the user email
@@ -355,4 +354,32 @@ def log_db_connection_error(connection_type: str, error_message: str) -> None:
             "error_message": error_message,
         }
     )
+    
+#################################
+# Streamlit Conversion
+#################################
 
+def log_streamlit_app_creation_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType) -> None:
+    log(
+        "mito_ai_streamlit_app_creation_success", 
+        key_type=key_type
+    )
+    
+def log_streamlit_app_creation_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str) -> None:
+    log(
+        "mito_ai_streamlit_app_creation_retry", 
+        params={
+            "error_message": error,
+        },
+        key_type=key_type
+    )
+    
+def log_streamlit_app_creation_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str) -> None:
+    log(
+        "mito_ai_streamlit_app_creation_error", 
+        params={
+            "error_message": error,
+        },
+        key_type=key_type
+    )
+    

{mito_ai-0.1.37.data → mito_ai-0.1.39.data}/data/share/jupyter/labextensions/mito_ai/build_log.json

@@ -710,7 +710,7 @@
             "semver": {},
             "vscode-diff": {},
             "mito_ai": {
-              "version": "0.1.37",
+              "version": "0.1.39",
               "singleton": true,
               "import": "/home/runner/work/mito/mito/mito-ai/lib/index.js"
             }

{mito_ai-0.1.37.data → mito_ai-0.1.39.data}/data/share/jupyter/labextensions/mito_ai/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.37",
+  "version": "0.1.39",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -138,7 +138,7 @@
     "outputDir": "mito_ai/labextension",
     "schemaDir": "schema",
     "_build": {
-      "load": "static/remoteEntry.93ecc9bc0edba61535cc.js",
+      "load": "static/remoteEntry.606207904e6aaa42b1bf.js",
       "extension": "./extension",
       "style": "./style"
     }

{mito_ai-0.1.37.data → mito_ai-0.1.39.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/package.json.orig

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.37",
+  "version": "0.1.39",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",

{mito_ai-0.1.37.data → mito_ai-0.1.39.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/toolbar-buttons.json

@@ -27,10 +27,15 @@
         }
       ],
       "Notebook": [
+        {
+          "name": "preview-app",
+          "command": "toolbar-button:preview-as-streamlit",
+          "rank": 1000
+        },
         {
           "name": "convert-to-streamlit",
           "command": "toolbar-button:convert-to-streamlit",
-          "rank": 1000
+          "rank": 1001
         }
       ]
     }