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

mito_ai/tests/streamlit_conversion/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -0,0 +1,265 @@
+# 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, AsyncMock
+from mito_ai.streamlit_conversion.streamlit_agent_handler import (
+    StreamlitCodeGeneration,
+    streamlit_handler
+)
+from typing import cast
+
+# Add this line to enable async support
+pytest_plugins = ('pytest_asyncio',)
+
+
+class TestStreamlitCodeGeneration:
+    """Test cases for StreamlitCodeGeneration class"""
+
+    def test_init(self):
+        """Test StreamlitCodeGeneration initialization"""
+        notebook_data: dict = {"cells": [{"cell_type": "code", "source": ["import pandas"]}]}
+        generator = StreamlitCodeGeneration(notebook_data)
+        
+        assert len(generator.messages) == 1
+        assert generator.messages[0]["role"] == "user"
+        # Access content properly as a list and cast to expected type
+        content_list = cast(list, generator.messages[0]["content"])
+        assert isinstance(content_list, list)
+        assert len(content_list) > 0
+        content_item = cast(dict, content_list[0])
+        assert content_item["type"] == "text"
+        assert "jupyter notebook content" in content_item["text"]
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.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(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()
+        
+        notebook_data: dict = {"cells": []}
+        generator = StreamlitCodeGeneration(notebook_data)
+        
+        result = await generator.get_response_from_agent(generator.messages)
+        
+        assert result == expected_result
+        mock_stream.assert_called_once()
+        
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.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")
+        
+        notebook_data: dict = {"cells": []}
+        generator = StreamlitCodeGeneration(notebook_data)
+        
+        with pytest.raises(Exception, match="API Error"):
+            await generator.get_response_from_agent(generator.messages)
+
+    def test_add_agent_response_to_context(self):
+        """Test adding agent response to message history"""
+        notebook_data: dict = {"cells": []}
+        generator = StreamlitCodeGeneration(notebook_data)
+        
+        initial_count = len(generator.messages)
+        generator.add_agent_response_to_context("Test response")
+        
+        assert len(generator.messages) == initial_count + 1
+        assert generator.messages[-1]["role"] == "assistant"
+        # Access content properly as a list and cast to expected type
+        content_list = cast(list, generator.messages[-1]["content"])
+        assert isinstance(content_list, list)
+        assert len(content_list) > 0
+        content_item = cast(dict, content_list[0])
+        assert content_item["type"] == "text"
+        assert content_item["text"] == "Test response"
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
+    async def test_generate_streamlit_code_success(self, mock_stream):
+        """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()
+        
+        notebook_data: dict = {"cells": []}
+        generator = StreamlitCodeGeneration(notebook_data)
+        
+        result = await generator.generate_streamlit_code()
+        
+        expected_code = "import streamlit\nst.title('Hello')\n"
+        assert result == expected_code
+        
+        # Check that response was added to context
+        assert len(generator.messages) == 2
+        assert generator.messages[-1]["role"] == "assistant"
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
+    async def test_correct_error_in_generation_success(self, mock_stream):
+        """Test successful error correction"""
+        mock_response = "Here's the corrected code:\n```python\nimport streamlit\nst.title('Fixed')\n```"
+        async def mock_async_gen():
+            for item in [mock_response]:
+                yield item
+
+        mock_stream.return_value = mock_async_gen()
+
+        notebook_data: dict = {"cells": []}
+        generator = StreamlitCodeGeneration(notebook_data)
+
+        result = await generator.correct_error_in_generation("ImportError: No module named 'pandas'")
+
+        expected_code = "import streamlit\nst.title('Fixed')\n"
+        assert result == expected_code
+
+        # Check that error message and response were added to context
+        assert len(generator.messages) == 3
+        assert generator.messages[-2]["role"] == "user"
+        assert generator.messages[-1]["role"] == "assistant"
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.stream_anthropic_completion_from_mito_server')
+    async def test_correct_error_in_generation_exception(self, mock_stream):
+        """Test exception handling in error correction"""
+        mock_stream.side_effect = Exception("API Error")
+        
+        notebook_data: dict = {"cells": []}
+        generator = StreamlitCodeGeneration(notebook_data)
+        
+        with pytest.raises(Exception, match="API Error"):
+            await generator.correct_error_in_generation("Some error")
+
+
+class TestStreamlitHandler:
+    """Test cases for streamlit_handler function"""
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.StreamlitCodeGeneration')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.streamlit_code_validator')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
+    async def test_streamlit_handler_success(self, mock_create_file, mock_validator, mock_generator_class, mock_parse):
+        """Test successful streamlit handler execution"""
+        # Mock notebook parsing
+        mock_notebook_data: dict = {"cells": [{"cell_type": "code", "source": ["import pandas"]}]}
+        mock_parse.return_value = mock_notebook_data
+        
+        # Mock code generation
+        mock_generator = AsyncMock()
+        mock_generator.generate_streamlit_code.return_value = "import streamlit\nst.title('Test')"
+        mock_generator_class.return_value = mock_generator
+        
+        # Mock validation (no errors)
+        mock_validator.return_value = (False, "")
+        
+        # Mock file creation
+        mock_create_file.return_value = (True, "File created successfully")
+        
+        result = await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+        
+        assert result[0] is True
+        assert "File created successfully" in result[1]
+        
+        # Verify calls
+        mock_parse.assert_called_once_with("/path/to/notebook.ipynb")
+        mock_generator_class.assert_called_once_with(mock_notebook_data)
+        mock_generator.generate_streamlit_code.assert_called_once()
+        mock_validator.assert_called_once_with("import streamlit\nst.title('Test')")
+        mock_create_file.assert_called_once_with("/path/to/app", "import streamlit\nst.title('Test')")
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.StreamlitCodeGeneration')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.streamlit_code_validator')
+    async def test_streamlit_handler_max_retries_exceeded(self, mock_validator, mock_generator_class, mock_parse):
+        """Test streamlit handler when max retries are exceeded"""
+        # Mock notebook parsing
+        mock_notebook_data: dict = {"cells": []}
+        mock_parse.return_value = mock_notebook_data
+        
+        # Mock code generation
+        mock_generator = AsyncMock()
+        mock_generator.generate_streamlit_code.return_value = "import streamlit\nst.title('Test')"
+        mock_generator.correct_error_in_generation.return_value = "import streamlit\nst.title('Fixed')"
+        mock_generator_class.return_value = mock_generator
+        
+        # Mock validation (always errors)
+        mock_validator.return_value = (True, "Persistent error")
+        
+        result = await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+        
+        assert result[0] is False
+        assert "Error generating streamlit code by agent" in result[1]
+        
+        # Verify that error correction was called 5 times (max retries)
+        assert mock_generator.correct_error_in_generation.call_count == 5
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.StreamlitCodeGeneration')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.streamlit_code_validator')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
+    async def test_streamlit_handler_file_creation_failure(self, mock_create_file, mock_validator, mock_generator_class, mock_parse):
+        """Test streamlit handler when file creation fails"""
+        # Mock notebook parsing
+        mock_notebook_data: dict = {"cells": []}
+        mock_parse.return_value = mock_notebook_data
+        
+        # Mock code generation
+        mock_generator = AsyncMock()
+        mock_generator.generate_streamlit_code.return_value = "import streamlit\nst.title('Test')"
+        mock_generator_class.return_value = mock_generator
+        
+        # Mock validation (no errors)
+        mock_validator.return_value = (False, "")
+        
+        # Mock file creation failure
+        mock_create_file.return_value = (False, "Permission denied")
+        
+        result = await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+        
+        assert result[0] is False
+        assert "Permission denied" in result[1]
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
+    async def test_streamlit_handler_parse_notebook_exception(self, mock_parse):
+        """Test streamlit handler when notebook parsing fails"""
+        mock_parse.side_effect = FileNotFoundError("Notebook not found")
+        
+        with pytest.raises(FileNotFoundError, match="Notebook not found"):
+            await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+
+    @pytest.mark.asyncio
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.StreamlitCodeGeneration')
+    async def test_streamlit_handler_generation_exception(self, mock_generator_class, mock_parse):
+        """Test streamlit handler when code generation fails"""
+        # Mock notebook parsing
+        mock_notebook_data: dict = {"cells": []}
+        mock_parse.return_value = mock_notebook_data
+        
+        # Mock code generation failure
+        mock_generator = AsyncMock()
+        mock_generator.generate_streamlit_code.side_effect = Exception("Generation failed")
+        mock_generator_class.return_value = mock_generator
+        
+        with pytest.raises(Exception, match="Generation failed"):
+            await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")

mito_ai/tests/streamlit_conversion/test_streamlit_utils.py

@@ -0,0 +1,197 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import json
+import tempfile
+import os
+from unittest.mock import patch, mock_open
+from mito_ai.streamlit_conversion.streamlit_utils import (
+    extract_code_blocks,
+    create_app_file,
+    parse_jupyter_notebook_to_extract_required_content
+)
+from typing import Dict, Any
+
+
+class TestExtractCodeBlocks:
+    """Test cases for extract_code_blocks function"""
+
+    def test_extract_code_blocks_with_python_blocks(self):
+        """Test extracting code from message with python code blocks"""
+        message = "Here's some code:\n```python\nimport streamlit\nst.title('Hello')\n```\nThat's it!"
+        result = extract_code_blocks(message)
+        expected = "import streamlit\nst.title('Hello')\n"
+        assert result == expected
+
+    def test_extract_code_blocks_without_python_blocks(self):
+        """Test when message doesn't contain python code blocks"""
+        message = "This is just regular text without code blocks"
+        result = extract_code_blocks(message)
+        assert result == message
+
+    def test_extract_code_blocks_empty_message(self):
+        """Test with empty message"""
+        message = ""
+        result = extract_code_blocks(message)
+        assert result == message
+
+    def test_extract_code_blocks_multiple_blocks(self):
+        """Test extracting from first python block when multiple exist"""
+        message = "```python\ncode1\n```\n```python\ncode2\n```"
+        result = extract_code_blocks(message)
+        expected = "code1\n\ncode2\n"
+        assert result == expected
+
+
+class TestCreateAppFile:
+    """Test cases for create_app_file function"""
+
+    def test_create_app_file_success(self, tmp_path):
+        """Test successful file creation"""
+        file_path = str(tmp_path)
+        code = "import streamlit\nst.title('Test')"
+        
+        success, message = create_app_file(file_path, code)
+        
+        assert success is True
+        assert "Successfully created" in message
+        
+        # Verify file was created with correct content
+        app_file_path = os.path.join(file_path, "app.py")
+        assert os.path.exists(app_file_path)
+        
+        with open(app_file_path, 'r') as f:
+            content = f.read()
+        assert content == code
+
+    def test_create_app_file_io_error(self):
+        """Test file creation with IO error"""
+        file_path = "/nonexistent/path/that/should/fail"
+        code = "import streamlit"
+        
+        success, message = create_app_file(file_path, code)
+        
+        assert success is False
+        assert "Error creating file" in message
+
+    @patch('builtins.open', side_effect=Exception("Unexpected error"))
+    def test_create_app_file_unexpected_error(self, mock_open):
+        """Test file creation with unexpected error"""
+        file_path = "/tmp/test"
+        code = "import streamlit"
+        
+        success, message = create_app_file(file_path, code)
+        
+        assert success is False
+        assert "Unexpected error" in message
+
+    def test_create_app_file_empty_code(self, tmp_path):
+        """Test creating file with empty code"""
+        file_path = str(tmp_path)
+        code = ""
+        
+        success, message = create_app_file(file_path, code)
+        
+        assert success is True
+        assert "Successfully created" in message
+        
+        app_file_path = os.path.join(file_path, "app.py")
+        with open(app_file_path, 'r') as f:
+            content = f.read()
+        assert content == ""
+
+
+class TestParseJupyterNotebookToExtractRequiredContent:
+    """Test cases for parse_jupyter_notebook_to_extract_required_content function"""
+
+    def test_parse_valid_notebook(self, tmp_path):
+        """Test parsing a valid notebook with cells"""
+        notebook_data: Dict[str, Any] = {
+            "cells": [
+                {
+                    "cell_type": "code",
+                    "source": ["import pandas as pd\n", "df = pd.DataFrame()\n"],
+                    "metadata": {"some": "metadata"},
+                    "execution_count": 1
+                },
+                {
+                    "cell_type": "markdown",
+                    "source": ["# Title\n", "Some text\n"],
+                    "metadata": {"another": "metadata"}
+                }
+            ],
+            "metadata": {"notebook_metadata": "value"},
+            "nbformat": 4
+        }
+        
+        notebook_path = tmp_path / "test.ipynb"
+        with open(notebook_path, 'w') as f:
+            json.dump(notebook_data, f)
+        
+        result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
+        
+        # Check that only cell_type and source are preserved
+        assert len(result['cells']) == 2
+        assert result['cells'][0]['cell_type'] == 'code'
+        assert result['cells'][0]['source'] == ["import pandas as pd\n", "df = pd.DataFrame()\n"]
+        assert 'metadata' not in result['cells'][0]
+        assert 'execution_count' not in result['cells'][0]
+        
+        assert result['cells'][1]['cell_type'] == 'markdown'
+        assert result['cells'][1]['source'] == ["# Title\n", "Some text\n"]
+        assert 'metadata' not in result['cells'][1]
+
+    def test_parse_notebook_file_not_found(self):
+        """Test parsing non-existent notebook file"""
+        with pytest.raises(FileNotFoundError, match="Notebook file not found"):
+            parse_jupyter_notebook_to_extract_required_content("/nonexistent/path/notebook.ipynb")
+
+    def test_parse_notebook_with_missing_cell_fields(self, tmp_path):
+        """Test parsing notebook where cells are missing cell_type or source"""
+        notebook_data: Dict[str, Any] = {
+            "cells": [
+                {
+                    "cell_type": "code"
+                    # Missing source field
+                },
+                {
+                    "source": ["some text"]
+                    # Missing cell_type field
+                },
+                {
+                    "cell_type": "markdown",
+                    "source": ["# Title"]
+                }
+            ]
+        }
+        
+        notebook_path = tmp_path / "test.ipynb"
+        with open(notebook_path, 'w') as f:
+            json.dump(notebook_data, f)
+        
+        result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
+        
+        assert len(result['cells']) == 3
+        assert result['cells'][0]['cell_type'] == 'code'
+        assert result['cells'][0]['source'] == []  # Default empty list
+        
+        assert result['cells'][1]['cell_type'] == ''  # Default empty string
+        assert result['cells'][1]['source'] == ["some text"]
+        
+        assert result['cells'][2]['cell_type'] == 'markdown'
+        assert result['cells'][2]['source'] == ["# Title"]
+
+    def test_parse_empty_notebook(self, tmp_path):
+        """Test parsing notebook with empty cells list"""
+        notebook_data: Dict[str, Any] = {
+            "cells": []
+        }
+        
+        notebook_path = tmp_path / "test.ipynb"
+        with open(notebook_path, 'w') as f:
+            json.dump(notebook_data, f)
+        
+        result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
+        
+        assert result['cells'] == []