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

mito_ai/tests/streamlit_conversion/test_validate_and_run_streamlit_code.py

@@ -0,0 +1,418 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import tempfile
+import os
+import subprocess
+import time
+import ast
+import importlib.util
+from unittest.mock import patch, MagicMock, mock_open
+from mito_ai.streamlit_conversion.validate_and_run_streamlit_code import (
+    StreamlitValidator,
+    streamlit_code_validator
+)
+
+
+class TestStreamlitValidator:
+    """Test cases for StreamlitValidator class"""
+
+    def test_validate_syntax_valid_code(self):
+        """Test syntax validation with valid Python code"""
+        validator = StreamlitValidator()
+        code = "import streamlit\nst.title('Hello World')"
+        
+        is_valid, message = validator.validate_syntax(code)
+        
+        assert is_valid is True
+        assert "Syntax is valid" in message
+
+    def test_validate_syntax_invalid_code(self):
+        """Test syntax validation with invalid Python code"""
+        validator = StreamlitValidator()
+        code = "import streamlit\nst.title('Hello World'  # Missing closing parenthesis"
+        
+        is_valid, message = validator.validate_syntax(code)
+        
+        assert is_valid is False
+        assert "Syntax error" in message
+
+    def test_validate_syntax_empty_code(self):
+        """Test syntax validation with empty code"""
+        validator = StreamlitValidator()
+        code = ""
+        
+        is_valid, message = validator.validate_syntax(code)
+        
+        assert is_valid is True
+        assert "Syntax is valid" in message
+
+    def test_create_temp_app(self):
+        """Test creating temporary app file"""
+        validator = StreamlitValidator()
+        code = "import streamlit\nst.title('Test')"
+        
+        app_path = validator.create_temp_app(code)
+        
+        assert validator.temp_dir is not None
+        assert os.path.exists(validator.temp_dir)
+        assert app_path.endswith("app.py")
+        assert os.path.exists(app_path)
+        
+        # Check file content
+        with open(app_path, 'r') as f:
+            content = f.read()
+        assert content == code
+        
+        # Cleanup
+        validator.cleanup()
+
+    def test_create_temp_app_empty_code(self):
+        """Test creating temporary app file with empty code"""
+        validator = StreamlitValidator()
+        code = ""
+        
+        app_path = validator.create_temp_app(code)
+        
+        assert os.path.exists(app_path)
+        
+        with open(app_path, 'r') as f:
+            content = f.read()
+        assert content == ""
+        
+        # Cleanup
+        validator.cleanup()
+
+    @patch('subprocess.Popen')
+    def test_start_streamlit_app_success(self, mock_popen):
+        """Test successful Streamlit app startup"""
+        validator = StreamlitValidator(port=8502)
+        app_path = "/tmp/test/app.py"
+        
+        # Mock successful subprocess
+        mock_process = MagicMock()
+        mock_popen.return_value = mock_process
+        
+        success, message = validator.start_streamlit_app(app_path)
+        
+        assert success is True
+        assert "Streamlit app started" in message
+        assert validator.process == mock_process
+        
+        # Verify subprocess was called with correct arguments
+        mock_popen.assert_called_once()
+        call_args = mock_popen.call_args[0][0]
+        assert "streamlit" in call_args
+        assert "run" in call_args
+        assert app_path in call_args
+        assert "8502" in call_args
+
+    @patch('subprocess.Popen')
+    def test_start_streamlit_app_failure(self, mock_popen):
+        """Test Streamlit app startup failure"""
+        validator = StreamlitValidator()
+        app_path = "/tmp/test/app.py"
+        
+        # Mock subprocess failure
+        mock_popen.side_effect = Exception("Failed to start process")
+        
+        success, message = validator.start_streamlit_app(app_path)
+        
+        assert success is False
+        assert "Failed to start Streamlit" in message
+        assert validator.process is None
+
+    @patch('requests.get')
+    def test_wait_for_app_success(self, mock_get):
+        """Test waiting for app to be ready successfully"""
+        validator = StreamlitValidator(port=8502, timeout=5)
+        
+        # Mock successful HTTP response
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_get.return_value = mock_response
+        
+        success, message = validator.wait_for_app()
+        
+        assert success is True
+        assert "App is running successfully" in message
+        mock_get.assert_called_with("http://localhost:8502", timeout=5)
+
+    @patch('requests.get')
+    def test_wait_for_app_http_error(self, mock_get):
+        """Test waiting for app with HTTP error"""
+        validator = StreamlitValidator(port=8501, timeout=5)
+        
+        # Mock HTTP error response
+        mock_response = MagicMock()
+        mock_response.status_code = 500
+        mock_get.return_value = mock_response
+        
+        success, message = validator.wait_for_app()
+        
+        assert success is False
+        assert "App failed to start within timeout" in message
+
+    @patch('subprocess.Popen')
+    def test_check_for_errors_process_running(self, mock_popen):
+        """Test error checking when process is running"""
+        validator = StreamlitValidator()
+        
+        # Mock running process
+        mock_process = MagicMock()
+        mock_process.poll.return_value = None  # Process is running
+        validator.process = mock_process
+        
+        success, message = validator.check_for_errors()
+        
+        assert success is True
+        assert "App is running without errors" in message
+
+    @patch('subprocess.Popen')
+    def test_check_for_errors_process_crashed(self, mock_popen):
+        """Test error checking when process has crashed"""
+        validator = StreamlitValidator()
+        
+        # Mock crashed process
+        mock_process = MagicMock()
+        mock_process.poll.return_value = 1  # Process has exited
+        mock_process.communicate.return_value = ("stdout", "stderr error message")
+        validator.process = mock_process
+        
+        success, message = validator.check_for_errors()
+        
+        assert success is False
+        assert "App crashed" in message
+        assert "stderr error message" in message
+
+    @patch('subprocess.Popen')
+    def test_check_for_errors_process_crashed_with_warnings(self, mock_popen):
+        """Test error checking when process crashed but only has warnings"""
+        validator = StreamlitValidator()
+        
+        # Mock crashed process with only warnings
+        mock_process = MagicMock()
+        mock_process.poll.return_value = 1
+        mock_process.communicate.return_value = ("stdout", "missing ScriptRunContext warning")
+        validator.process = mock_process
+        
+        success, message = validator.check_for_errors()
+        
+        assert success is True
+        assert "App is running without errors" in message
+
+    def test_check_for_errors_no_process(self):
+        """Test error checking when no process exists"""
+        validator = StreamlitValidator()
+        
+        success, message = validator.check_for_errors()
+        
+        assert success is False
+        assert "No process found" in message
+
+    @patch('subprocess.Popen')
+    def test_cleanup_with_process(self, mock_popen):
+        """Test cleanup with running process"""
+        validator = StreamlitValidator()
+        
+        # Mock process
+        mock_process = MagicMock()
+        validator.process = mock_process
+        validator.temp_dir = "/tmp/test_dir"
+        
+        # Mock directory exists
+        with patch('os.path.exists', return_value=True):
+            with patch('shutil.rmtree') as mock_rmtree:
+                validator.cleanup()
+                
+                mock_process.terminate.assert_called_once()
+                mock_process.wait.assert_called_once()
+                mock_rmtree.assert_called_once_with("/tmp/test_dir")
+                
+                assert validator.process is None
+                assert validator.temp_dir is None  # type: ignore[unreachable]
+
+    def test_cleanup_without_process(self):
+        """Test cleanup without process"""
+        validator = StreamlitValidator()
+        
+        # Should not raise any exceptions
+        validator.cleanup()
+        
+        assert validator.process is None
+        assert validator.temp_dir is None
+
+    @patch('subprocess.Popen')
+    @patch('requests.get')
+    def test_validate_app_success(self, mock_get, mock_popen):
+        """Test complete validation pipeline success"""
+        validator = StreamlitValidator(port=8501, timeout=5)
+        code = "import streamlit\nst.title('Hello World')"
+        
+        # Mock successful subprocess
+        mock_process = MagicMock()
+        mock_process.poll.return_value = None
+        mock_popen.return_value = mock_process
+        
+        # Mock successful HTTP response
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_get.return_value = mock_response
+        
+        results = validator.validate_app(code)
+        
+        assert results['syntax_valid'] is True
+        assert results['app_starts'] is True
+        assert results['app_responsive'] is True
+        assert len(results['errors']) == 0
+
+    def test_validate_app_syntax_error(self):
+        """Test validation pipeline with syntax error"""
+        validator = StreamlitValidator()
+        code = "import streamlit\nst.title('Hello World'  # Missing parenthesis"
+        
+        results = validator.validate_app(code)
+        
+        assert results['syntax_valid'] is False
+        assert results['app_starts'] is False
+        assert results['app_responsive'] is False
+        assert len(results['errors']) == 1
+        assert "Syntax error" in results['errors'][0]
+
+    @patch('subprocess.Popen')
+    def test_validate_app_startup_failure(self, mock_popen):
+        """Test validation pipeline with startup failure"""
+        validator = StreamlitValidator()
+        code = "import streamlit\nst.title('Hello World')"
+        
+        # Mock startup failure
+        mock_popen.side_effect = Exception("Failed to start")
+        
+        results = validator.validate_app(code)
+        
+        assert results['syntax_valid'] is True
+        assert results['app_starts'] is False
+        assert results['app_responsive'] is False
+        assert len(results['errors']) == 1
+        assert "Failed to start Streamlit" in results['errors'][0]
+
+    @patch('subprocess.Popen')
+    @patch('requests.get')
+    def test_validate_app_runtime_error(self, mock_get, mock_popen):
+        """Test validation pipeline with runtime error"""
+        validator = StreamlitValidator(port=8501, timeout=5)
+        code = "import streamlit\nst.title('Hello World')"
+        
+        # Mock successful startup
+        mock_process = MagicMock()
+        mock_process.poll.return_value = 1  # Process crashed
+        mock_process.communicate.return_value = ("stdout", "Runtime error occurred")
+        mock_popen.return_value = mock_process
+        
+        # Mock successful HTTP response
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_get.return_value = mock_response
+        
+        results = validator.validate_app(code)
+        
+        assert results['syntax_valid'] is True
+        assert results['app_starts'] is True
+        assert results['app_responsive'] is True
+        assert len(results['errors']) == 1
+        assert "App crashed" in results['errors'][0]
+
+    def test_validate_app_exception_handling(self):
+        """Test validation pipeline exception handling"""
+        validator = StreamlitValidator()
+        code = "import streamlit\nst.title('Hello World')"
+        
+        # Mock an exception during validation
+        with patch.object(validator, 'validate_syntax', side_effect=Exception("Unexpected error")):
+            results = validator.validate_app(code)
+            
+            assert results['syntax_valid'] is False
+            assert results['app_starts'] is False
+            assert results['app_responsive'] is False
+            assert len(results['errors']) == 1
+            assert "Validation error" in results['errors'][0]
+
+
+class TestStreamlitCodeValidator:
+    """Test cases for streamlit_code_validator function"""
+
+    def test_streamlit_code_validator_success(self):
+        """Test successful code validation"""
+        code = "import streamlit\nst.title('Hello World')"
+        
+        with patch('mito_ai.streamlit_conversion.validate_and_run_streamlit_code.StreamlitValidator') as mock_validator_class:
+            mock_validator = MagicMock()
+            mock_validator_class.return_value = mock_validator
+            
+            mock_validator.validate_app.return_value = {
+                'syntax_valid': True,
+                'app_starts': True,
+                'app_responsive': True,
+                'errors': []
+            }
+            
+            has_error, message = streamlit_code_validator(code)
+            
+            assert has_error is False
+            assert "Errors found" not in message
+            mock_validator.validate_app.assert_called_once_with(code)
+
+    def test_streamlit_code_validator_error_in_code(self):
+        """Test code validation when code contains 'error'"""
+        code = "error in the code"
+        
+        has_error, message = streamlit_code_validator(code)
+        
+        assert has_error is True
+        assert "Errors found" in message
+
+    def test_streamlit_code_validator_empty_code(self):
+        """Test code validation with empty code"""
+        code = ""
+        
+        with patch('mito_ai.streamlit_conversion.validate_and_run_streamlit_code.StreamlitValidator') as mock_validator_class:
+            mock_validator = MagicMock()
+            mock_validator_class.return_value = mock_validator
+            
+            mock_validator.validate_app.return_value = {
+                'syntax_valid': True,
+                'app_starts': True,
+                'app_responsive': True,
+                'errors': []
+            }
+            
+            has_error, message = streamlit_code_validator(code)
+            
+            assert has_error is False
+            assert "Errors found" not in message
+
+    def test_streamlit_code_validator_multiple_errors(self):
+        """Test code validation with multiple errors"""
+        code = "import streamlit\nst.title('Hello World')"
+        
+        with patch('mito_ai.streamlit_conversion.validate_and_run_streamlit_code.StreamlitValidator') as mock_validator_class:
+            mock_validator = MagicMock()
+            mock_validator_class.return_value = mock_validator
+            
+            mock_validator.validate_app.return_value = {
+                'syntax_valid': False,
+                'app_starts': False,
+                'app_responsive': False,
+                'errors': [
+                    'Syntax error: invalid syntax',
+                    'App failed to start'
+                ]
+            }
+            
+            has_error, message = streamlit_code_validator(code)
+            
+            assert has_error is True
+            assert "Errors found" in message
+            assert "Syntax error" in message
+            assert "App failed to start" in message

mito_ai/tests/test_constants.py

@@ -3,8 +3,11 @@
 
 from typing import Any
 import pytest
-from mito_ai.constants import ACTIVE_BASE_URL, MITO_PROD_BASE_URL, MITO_DEV_BASE_URL
-from mito_ai.constants import MITO_STREAMLIT_DEV_BASE_URL, MITO_STREAMLIT_TEST_BASE_URL, ACTIVE_STREAMLIT_BASE_URL
+from mito_ai.constants import (
+    ACTIVE_BASE_URL, MITO_PROD_BASE_URL, MITO_DEV_BASE_URL,
+    MITO_STREAMLIT_DEV_BASE_URL, MITO_STREAMLIT_TEST_BASE_URL, ACTIVE_STREAMLIT_BASE_URL,
+    COGNITO_CONFIG_DEV, ACTIVE_COGNITO_CONFIG,
+)
 
 
 def test_prod_lambda_url() -> Any:
@@ -29,4 +32,16 @@ def test_testenv_streamlit_url() -> Any:
 
 def test_streamlit_active_base_url() -> Any:
     """Make sure that the active streamlit base url is correct"""
-    assert ACTIVE_STREAMLIT_BASE_URL == MITO_STREAMLIT_TEST_BASE_URL
+    assert ACTIVE_STREAMLIT_BASE_URL == MITO_STREAMLIT_DEV_BASE_URL
+
+def test_cognito_config() -> Any:
+    """Make sure that the Cognito configuration is correct"""
+    expected_config = {
+        'TOKEN_ENDPOINT': 'https://mito-app-auth.auth.us-east-1.amazoncognito.com/oauth2/token',
+        'CLIENT_ID': '6ara3u3l8sss738hrhbq1qtiqf',
+        'CLIENT_SECRET': '',
+        'REDIRECT_URI': 'http://localhost:8888/lab'
+    }
+
+    assert COGNITO_CONFIG_DEV == expected_config
+    assert ACTIVE_COGNITO_CONFIG == COGNITO_CONFIG_DEV

mito_ai/utils/anthropic_utils.py

@@ -8,7 +8,7 @@ import anthropic
 from typing import Any, Dict, List, Optional, Union, AsyncGenerator, Tuple, Callable, cast
 
 from anthropic.types import MessageParam, Message, TextBlock, ToolUnionParam
-from mito_ai.utils.mito_server_utils import get_response_from_mito_server
+from mito_ai.utils.mito_server_utils import get_response_from_mito_server, stream_response_from_mito_server
 from mito_ai.utils.provider_utils import does_message_require_fast_model
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import AgentResponse, MessageType, ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem
@@ -110,75 +110,23 @@ async def stream_anthropic_completion_from_mito_server(
     data, headers = _prepare_anthropic_request_data_and_headers(
         model, max_tokens, temperature, system, messages, message_type, None, None, stream
     )
-    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
-    start_time = time.time()
-    chunk_queue: asyncio.Queue[str] = asyncio.Queue()
-    fetch_complete = False
-    def chunk_callback(chunk: bytes) -> None:
-        try:
-            chunk_str = chunk.decode('utf-8')
-            asyncio.create_task(chunk_queue.put(chunk_str))
-        except Exception as e:
-            print(f"Error processing Anthropic streaming chunk: {str(e)}")
-    fetch_future = None
-    try:
-        fetch_future = http_client.fetch(
-            MITO_ANTHROPIC_URL,
-            method="POST",
-            headers=headers,
-            body=json.dumps(data),
-            request_timeout=http_client_timeout,
-            streaming_callback=chunk_callback
-        )
-        async def wait_for_fetch() -> None:
-            try:
-                await fetch_future
-                nonlocal fetch_complete
-                fetch_complete = True
-                print("Anthropic fetch completed")
-            except Exception as e:
-                print(f"Error in Anthropic fetch: {str(e)}")
-                raise
-        fetch_task = asyncio.create_task(wait_for_fetch())
-        while not (fetch_complete and chunk_queue.empty()):
-            try:
-                chunk = await asyncio.wait_for(chunk_queue.get(), timeout=0.1)
-                if reply_fn and message_id:
-                    reply_fn(CompletionStreamChunk(
-                        parent_id=message_id,
-                        chunk=CompletionItem(
-                            content=chunk,
-                            isIncomplete=True,
-                            token=message_id,
-                        ),
-                        done=False,
-                    ))
-                yield chunk
-            except asyncio.TimeoutError:
-                if fetch_complete and chunk_queue.empty():
-                    break
-                continue
-        print(f"\nAnthropic stream completed in {time.time() - start_time:.2f} seconds")
-        if reply_fn and message_id:
-            reply_fn(CompletionStreamChunk(
-                parent_id=message_id,
-                chunk=CompletionItem(
-                    content="",
-                    isIncomplete=False,
-                    token=message_id,
-                ),
-                done=True,
-            ))
-    except Exception as e:
-        print(f"\nAnthropic stream failed after {time.time() - start_time:.2f} seconds with error: {str(e)}")
-        if fetch_future:
-            try:
-                await fetch_future
-            except Exception:
-                pass
-        raise
-    finally:
-        http_client.close()
+    # Use the unified streaming function
+    # If the reply_fn and message_id are empty, this function still handles those requests. This is particularly needed for the streamlit dashboard functionality
+    actual_reply_fn = reply_fn if reply_fn is not None else (lambda x: None)
+    actual_message_id = message_id if message_id is not None else ""
+    async for chunk in stream_response_from_mito_server(
+        url=MITO_ANTHROPIC_URL,
+        headers=headers,
+        data=data,
+        timeout=timeout,
+        max_retries=max_retries,
+        message_type=message_type,
+        reply_fn=actual_reply_fn,
+        message_id=actual_message_id,
+        chunk_processor=None,
+        provider_name="Claude",
+    ):
+        yield chunk
 
 def get_anthropic_completion_function_params(
     message_type: MessageType,

mito_ai/utils/gemini_utils.py

@@ -5,7 +5,7 @@ import asyncio
 import json
 import time
 from typing import Any, Dict, List, Optional, Callable, Union, AsyncGenerator, Tuple
-from mito_ai.utils.mito_server_utils import get_response_from_mito_server
+from mito_ai.utils.mito_server_utils import get_response_from_mito_server, stream_response_from_mito_server
 from mito_ai.completions.models import AgentResponse, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
 from mito_ai.constants import MITO_GEMINI_URL
 from mito_ai.utils.provider_utils import does_message_require_fast_model
@@ -80,80 +80,29 @@ async def stream_gemini_completion_from_mito_server(
     contents: List[Dict[str, Any]],
     message_type: MessageType,
     message_id: str,
-    reply_fn: Optional[Callable[[Union[CompletionReply, CompletionStreamChunk]], None]]
+    reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]
 ) -> AsyncGenerator[str, None]:
     data, headers = _prepare_gemini_request_data_and_headers(model, contents, message_type, stream=True)
-    http_client, http_client_timeout = _create_http_client(timeout, max_retries)
-    start_time = time.time()
-    chunk_queue: asyncio.Queue[str] = asyncio.Queue()
-    fetch_complete = False
-    def chunk_callback(chunk: bytes) -> None:
-        try:
-            chunk_str = chunk.decode('utf-8')
-            asyncio.create_task(chunk_queue.put(chunk_str))
-        except Exception as e:
-            print(f"Error processing Gemini streaming chunk: {str(e)}")
-    fetch_future = None
-    try:
-        fetch_future = http_client.fetch(
-            MITO_GEMINI_URL,
-            method="POST",
-            headers=headers,
-            body=json.dumps(data),
-            request_timeout=http_client_timeout,
-            streaming_callback=chunk_callback
-        )
-        async def wait_for_fetch() -> None:
-            try:
-                await fetch_future
-                nonlocal fetch_complete
-                fetch_complete = True
-                print("Gemini fetch completed")
-            except Exception as e:
-                print(f"Error in Gemini fetch: {str(e)}")
-                raise
-        fetch_task = asyncio.create_task(wait_for_fetch())
-        while not (fetch_complete and chunk_queue.empty()):
-            try:
-                chunk = await asyncio.wait_for(chunk_queue.get(), timeout=0.1)
-                clean_chunk = chunk.strip('"')
-                decoded_chunk = clean_chunk.encode().decode('unicode_escape')
-                if reply_fn and message_id:
-                    reply_fn(CompletionStreamChunk(
-                        parent_id=message_id,
-                        chunk=CompletionItem(
-                            content=decoded_chunk,
-                            isIncomplete=True,
-                            token=message_id,
-                        ),
-                        done=False,
-                    ))
-                yield chunk
-            except asyncio.TimeoutError:
-                if fetch_complete and chunk_queue.empty():
-                    break
-                continue
-        print(f"\nGemini stream completed in {time.time() - start_time:.2f} seconds")
-        if reply_fn and message_id:
-            reply_fn(CompletionStreamChunk(
-                parent_id=message_id,
-                chunk=CompletionItem(
-                    content="",
-                    isIncomplete=False,
-                    token=message_id,
-                ),
-                done=True,
-            ))
-    except Exception as e:
-        print(f"\nGemini stream failed after {time.time() - start_time:.2f} seconds with error: {str(e)}")
-        if fetch_future:
-            try:
-                await fetch_future
-            except Exception:
-                pass
-        raise
-    finally:
-        http_client.close()
+    
+    # Define chunk processor for Gemini's special processing
+    def gemini_chunk_processor(chunk: str) -> str:
+        clean_chunk = chunk.strip('"')
+        return clean_chunk.encode().decode('unicode_escape')
+    
+    # Use the unified streaming function with Gemini's chunk processor
+    async for chunk in stream_response_from_mito_server(
+        url=MITO_GEMINI_URL,
+        headers=headers,
+        data=data,
+        timeout=timeout,
+        max_retries=max_retries,
+        message_type=message_type,
+        reply_fn=reply_fn,
+        message_id=message_id,
+        chunk_processor=gemini_chunk_processor,
+        provider_name="Gemini",
+    ):
+        yield chunk
 
 def get_gemini_completion_function_params(
     message_type: MessageType,