mito-ai 0.1.43__py3-none-any.whl → 0.1.45__py3-none-any.whl

mito_ai/completions/providers.py

@@ -160,7 +160,7 @@ This attribute is observed by the websocket provider to push the error to the cl
                 # If we hit a free tier limit, then raise an exception right away without retrying.
                 self.log.exception(f"Error during request_completions: {e}")
                 self.last_error = CompletionError.from_exception(e)
-                log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', message_type, e)
+                log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
                 raise
 
             except BaseException as e:
@@ -169,14 +169,14 @@ This attribute is observed by the websocket provider to push the error to the cl
                     # Exponential backoff: wait 2^attempt seconds
                     wait_time = 2 ** attempt
                     self.log.info(f"Retrying request_completions after {wait_time}s (attempt {attempt + 1}/{max_retries + 1}): {str(e)}")
-                    log_ai_completion_retry('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', message_type, e)
+                    log_ai_completion_retry('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
                     await asyncio.sleep(wait_time)
                     continue
                 else:
                     # Final failure after all retries - set error state and raise
                     self.log.exception(f"Error during request_completions after {attempt + 1} attempts: {e}")
                     self.last_error = CompletionError.from_exception(e)
-                    log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', message_type, e)
+                    log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
                     raise
         
         # This should never be reached due to the raise in the except block,
@@ -264,14 +264,8 @@ This attribute is observed by the websocket provider to push the error to the cl
         except BaseException as e:
             self.log.exception(f"Error during stream_completions: {e}")
             self.last_error = CompletionError.from_exception(e)
-            log(
-                MITO_AI_COMPLETION_ERROR,
-                params={
-                    KEY_TYPE_PARAM: self.key_type,
-                    'message_type': message_type.value,
-                },
-                error=e
-            )
+            log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id, message_type, e)
+
             # Send error message to client before raising
             reply_fn(CompletionStreamChunk(
                 parent_id=message_id,

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -4,7 +4,7 @@
 import logging
 import os
 from anthropic.types import MessageParam
-from typing import List, Optional, Tuple, cast, Union
+from typing import List, Optional, Tuple, cast
 
 from mito_ai.logger import get_logger
 from mito_ai.streamlit_conversion.agent_utils import apply_patch_to_text, extract_todo_placeholders, fix_diff_headers
@@ -52,18 +52,19 @@ class StreamlitCodeGeneration:
     async def generate_streamlit_code(self, notebook: dict) -> str:
         """Send a query to the agent, get its response and parse the code"""
         
+        prompt_text = get_streamlit_app_creation_prompt(notebook)
+        
         messages: List[MessageParam] = [
             cast(MessageParam, {
                 "role": "user",
                 "content": [{
                     "type": "text",
-                    "text": get_streamlit_app_creation_prompt(notebook)
+                    "text": prompt_text
                 }]
             })
         ]
         
         agent_response = await self.get_response_from_agent(messages)
-
         converted_code = extract_code_blocks(agent_response)
         
         # Extract the TODOs from the agent's response
@@ -123,6 +124,7 @@ async def streamlit_handler(notebook_path: str) -> Tuple[bool, Optional[str], st
 
     notebook_code = parse_jupyter_notebook_to_extract_required_content(notebook_path)
     streamlit_code_generator = StreamlitCodeGeneration()
+
     streamlit_code = await streamlit_code_generator.generate_streamlit_code(notebook_code)
     
     has_validation_error, errors = validate_app(streamlit_code, notebook_path)
@@ -149,6 +151,7 @@ async def streamlit_handler(notebook_path: str) -> Tuple[bool, Optional[str], st
         absolute_notebook_path = os.path.join(os.getcwd(), notebook_path)
     
     app_directory = os.path.dirname(absolute_notebook_path)
+    
     success_flag, app_path, message = create_app_file(app_directory, streamlit_code)
     
     if not success_flag:

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -25,7 +25,8 @@ def extract_code_blocks(message_content: str) -> str:
     matches = re.findall(pattern, message_content, re.DOTALL)
 
     # Concatenate with single newlines
-    return '\n'.join(matches)
+    result = '\n'.join(matches)
+    return result
 
 def extract_unified_diff_blocks(message_content: str) -> str:
     """
@@ -53,14 +54,26 @@ def create_app_file(app_directory: str, code: str) -> Tuple[bool, str, str]:
     """
     try:
         app_path = os.path.join(app_directory, "app.py")
-        with open(app_path, 'w') as f:
+
+        with open(app_path, 'w', encoding='utf-8') as f:
             f.write(code)
+        
         return True, app_path, f"Successfully created {app_directory}"
     except IOError as e:
         return False, '', f"Error creating file: {str(e)}"
     except Exception as e:
         return False, '', f"Unexpected error: {str(e)}"
+    
 
+def get_app_path(app_directory: str) -> Optional[str]:
+    """
+    Check if the app.py file exists in the given directory.
+    """
+    app_path = os.path.join(app_directory, "app.py")
+    if not os.path.exists(app_path):
+        return None
+    return app_path
+    
 
 def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Dict[str, Any]:
     """
@@ -129,8 +142,3 @@ def clean_directory_check(notebook_path: str) -> None:
 
     if not dir_path.exists():
         raise ValueError(f"Directory does not exist: {dir_path}")
-
-    file_count = len([f for f in dir_path.iterdir() if f.is_file()])
-    if file_count > 10:
-        raise ValueError(
-            f"Too many files in directory: 10 allowed but {file_count} present. Create a new directory and retry")

mito_ai/streamlit_conversion/validate_streamlit_app.py

@@ -25,7 +25,7 @@ warnings.filterwarnings("ignore", message=".*bare mode.*")
 
 class StreamlitValidator:
     def __init__(self, port: int = 8501) -> None:
-        self.temp_dir: Optional[str] = None
+        pass
 
     def get_syntax_error(self, app_code: str) -> Optional[str]:
         """Check if the Python code has valid syntax"""
@@ -58,26 +58,40 @@ class StreamlitValidator:
                 os.chdir(original_cwd)
         
         with change_working_directory(directory):
-            app_test = AppTest.from_string(app_code, default_timeout=30)
-            app_test.run()
-            
-            # Check for exceptions
-            if app_test.exception:
-                errors = [{'type': 'exception', 'details': exc.value, 'message': exc.message, 'stack_trace': exc.stack_trace} for exc in app_test.exception]
-                return errors
-                    
-            # Check for error messages
-            if app_test.error:
-                errors = [{'type': 'error', 'details': err.value} for err in app_test.error]
-                return errors
-            
-            return None
+            # Create a temporary file that uses UTF-8 encoding so 
+            # we don't run into issues with non-ASCII characters on Windows.
+            # We use utf-8 encoding when writing the app.py file so this validation
+            # code mirrors the actual file. 
+
+            # Note: Since the AppTest.from_file tries to open the file, we need to first close the file
+            # by exiting the context manager and using the delete=False flag so that the file still exists.
+            # Windows can't open the same file twice at the same time. We cleanup at the end.
+            with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8") as f:
+                f.write(app_code)
+                temp_path = f.name
 
-    def cleanup(self) -> None:
-        """Clean up the temporary files"""
-        if self.temp_dir and os.path.exists(self.temp_dir):
-            shutil.rmtree(self.temp_dir)
-            self.temp_dir = None
+            try:
+                # Run Streamlit test from file with UTF-8 encoding
+                app_test = AppTest.from_file(temp_path, default_timeout=30)
+                app_test.run()
+                
+                # Check for exceptions
+                if app_test.exception:
+                    errors = [{'type': 'exception', 'details': exc.value, 'message': exc.message, 'stack_trace': exc.stack_trace} for exc in app_test.exception]
+                    return errors
+                        
+                # Check for error messages
+                if app_test.error:
+                    errors = [{'type': 'error', 'details': err.value} for err in app_test.error]
+                    return errors
+                
+                return None
+            finally:
+                # Clean up the temporary file
+                try:
+                    os.unlink(temp_path)
+                except OSError:
+                    pass  # File might already be deleted
 
     def _validate_app(self, app_code: str, app_path: str) -> List[Dict[str, Any]]:
         """Complete validation pipeline"""
@@ -91,17 +105,12 @@ class StreamlitValidator:
 
             runtime_errors = self.get_runtime_errors(app_code, app_path)
             
-            print('Found Runtime Errors', runtime_errors)
-            
             if runtime_errors:
                 errors.extend(runtime_errors)
             
         except Exception as e:
             errors.append({'type': 'validation', 'details': str(e)})
 
-        finally:
-            self.cleanup()
-
         return errors
 
 def validate_app(app_code: str, notebook_path: str) -> Tuple[bool, List[str]]:

mito_ai/streamlit_preview/handlers.py

@@ -4,161 +4,140 @@
 import os
 import tempfile
 import uuid
+from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
+from mito_ai.streamlit_preview.utils import ensure_app_exists, validate_request_body
 import tornado
 from jupyter_server.base.handlers import APIHandler
 from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
 from mito_ai.streamlit_preview.manager import get_preview_manager
 from mito_ai.utils.create import initialize_user
+from typing import Tuple, Optional
+
 
 
 class StreamlitPreviewHandler(APIHandler):
     """REST handler for streamlit preview operations."""
-    
+
     def initialize(self) -> None:
         """Initialize the handler."""
         self.preview_manager = get_preview_manager()
-    
+
     def _resolve_notebook_path(self, notebook_path: str) -> str:
         """
         Resolve the notebook path to an absolute path that can be found by the backend.
-        
+
         This method handles path resolution issues that can occur in different environments:
-        
+
         1. **Test Environment**: Playwright tests create temporary directories with complex
            paths like 'mitoai_ui_tests-app_builde-ab3a5-n-Test-Preview-as-Streamlit-chromium/'
            that the backend can't directly access.
-        
+
         2. **JupyterHub/Cloud Deployments**: In cloud environments, users may have notebooks
            in subdirectories that aren't immediately accessible from the server root.
-        
+
         3. **Docker Containers**: When running in containers, the working directory and
            file paths may not align with what the frontend reports.
-        
+
         4. **Multi-user Environments**: In enterprise deployments, users may have notebooks
            in user-specific directories that require path resolution.
-        
+
         The method tries multiple strategies:
         1. If the path is already absolute, return it as-is
         2. Try to resolve relative to the Jupyter server's root directory
         3. Search recursively through subdirectories for a file with the same name
         4. Return the original path if not found (will cause a clear error message)
-        
+
         Args:
             notebook_path (str): The notebook path from the frontend (may be relative or absolute)
-            
+
         Returns:
             str: The resolved absolute path to the notebook file
         """
         # If the path is already absolute, return it
         if os.path.isabs(notebook_path):
             return notebook_path
-        
+
         # Get the Jupyter server's root directory
-        server_root = self.settings.get('server_root_dir', os.getcwd())
-        
+        server_root = self.settings.get("server_root_dir", os.getcwd())
+
         # Try to find the notebook file in the server root
         resolved_path = os.path.join(server_root, notebook_path)
         if os.path.exists(resolved_path):
             return resolved_path
-        
+
         # If not found, try to find it in subdirectories
         # This handles cases where the notebook is in a subdirectory that the frontend
         # doesn't know about, or where the path structure differs between frontend and backend
         for root, dirs, files in os.walk(server_root):
             if os.path.basename(notebook_path) in files:
                 return os.path.join(root, os.path.basename(notebook_path))
-        
+
         # If still not found, return the original path (will cause a clear error)
         # This ensures we get a meaningful error message rather than a generic "file not found"
         return os.path.join(os.getcwd(), notebook_path)
 
     @tornado.web.authenticated
     async def post(self) -> None:
-        """Start a new streamlit preview.
-        
-        Expected JSON body:
-        {
-            "notebook_path": "path/to/notebook.ipynb"
-        }
-        
-        Returns:
-        {
-            "id": "preview_id",
-            "port": 8501,
-            "url": "http://localhost:8501"
-        }
-        """
+        """Start a new streamlit preview."""
         try:
-            # Parse request body
+            # Parse and validate request
             body = self.get_json_body()
-            if body is None:
+            is_valid, error_msg, notebook_path, force_recreate = validate_request_body(body)
+            if not is_valid or notebook_path is None:
                 self.set_status(400)
-                self.finish({"error": 'Invalid or missing JSON body'})
+                self.finish({"error": error_msg})
                 return
 
-            notebook_path = body.get('notebook_path')
 
-            if not notebook_path:
-                self.set_status(400)
-                self.finish({"error": 'Missing notebook_path parameter'})
-                return
-            
-            # Resolve the notebook path to find the actual file
+            # Ensure app exists
             resolved_notebook_path = self._resolve_notebook_path(notebook_path)
-            
-            # Generate preview ID
-            preview_id = str(uuid.uuid4())
-            
-            # Generate streamlit code using existing handler
-            print('notebook_path', notebook_path)
-            success, app_path, message = await streamlit_handler(resolved_notebook_path)
-            
-            if not success or app_path is None:
+
+            success, error_msg = await ensure_app_exists(resolved_notebook_path, force_recreate)
+
+            if not success:
                 self.set_status(500)
-                self.finish({"error": f'Failed to generate streamlit code: {message}'})
+                self.finish({"error": error_msg})
                 return
-            
-            # Start streamlit preview
+
+            # Start preview
+            # TODO: There's a bug here where when the user rebuilds and already running app. Instead of 
+            # creating a new process, we should update the existing process. The app displayed to the user 
+            # does update, but that's just because of hot reloading when we overwrite the app.py file. 
+            preview_id = str(uuid.uuid4())
             resolved_app_directory = os.path.dirname(resolved_notebook_path)
             success, message, port = self.preview_manager.start_streamlit_preview(resolved_app_directory, preview_id)
-            
+
             if not success:
                 self.set_status(500)
-                self.finish({"error": f'Failed to start preview: {message}'})
+                self.finish({"error": f"Failed to start preview: {message}"})
                 return
-            
-            # Return success response - APIHandler automatically handles JSON serialization
-            self.finish({
-                'id': preview_id,
-                'port': port,
-                'url': f'http://localhost:{port}'
-            })
-                
+
+            # Return success response
+            self.finish({"id": preview_id, "port": port, "url": f"http://localhost:{port}"})
+
         except Exception as e:
             print(f"Error in streamlit preview handler: {e}")
             self.set_status(500)
-            
-            # Respond with the error
             self.finish({"error": str(e)})
-    
+
     @tornado.web.authenticated
     def delete(self, preview_id: str) -> None:
         """Stop a streamlit preview."""
         try:
             if not preview_id:
                 self.set_status(400)
-                self.finish({"error": 'Missing preview_id parameter'})
+                self.finish({"error": "Missing preview_id parameter"})
                 return
-            
+
             # Stop the preview
             stopped = self.preview_manager.stop_preview(preview_id)
-            
+
             if stopped:
                 self.set_status(204)  # No content
             else:
                 self.set_status(404)
-                self.finish({"error": f'Preview {preview_id} not found'})
-                
+                self.finish({"error": f"Preview {preview_id} not found"})
+
         except Exception as e:
             self.set_status(500)
-            self.finish({"error": f'Internal server error: {str(e)}'})
+            self.finish({"error": f"Internal server error: {str(e)}"})

mito_ai/streamlit_preview/utils.py

@@ -0,0 +1,41 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Tuple, Optional
+import os
+from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
+from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
+
+
+def validate_request_body(body: Optional[dict]) -> Tuple[bool, str, Optional[str], bool]:
+    """Validate the request body and extract notebook_path and force_recreate."""
+    if body is None:
+        return False, "Invalid or missing JSON body", None, False
+
+    notebook_path = body.get("notebook_path")
+    if not notebook_path:
+        return False, "Missing notebook_path parameter", None, False
+
+    force_recreate = body.get("force_recreate", False)
+    if not isinstance(force_recreate, bool):
+        return False, "force_recreate must be a boolean", None, False
+
+    return True, "", notebook_path, force_recreate
+
+async def ensure_app_exists(resolved_notebook_path: str, force_recreate: bool = False) -> Tuple[bool, str]:
+    """Ensure app.py exists, generating it if necessary or if force_recreate is True."""
+    # Check if the app already exists
+    app_path = get_app_path(os.path.dirname(resolved_notebook_path))
+    
+    if app_path is None or force_recreate:
+        if app_path is None:
+            print("[Mito AI] App path not found, generating streamlit code")
+        else:
+            print("[Mito AI] Force recreating streamlit app")
+        
+        success, app_path, message = await streamlit_handler(resolved_notebook_path)
+
+        if not success or app_path is None:
+            return False, f"Failed to generate streamlit code: {message}"
+
+    return True, ""

mito_ai/tests/deploy_app/test_app_deploy_utils.py

@@ -0,0 +1,71 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import zipfile
+import logging
+from mito_ai.app_deploy.app_deploy_utils import add_files_to_zip
+
+class TestAddFilesToZip:
+    """Test cases for add_files_to_zip helper function"""
+
+    def test_files_added_correctly(self, tmp_path):
+        """Ensure individual files are added correctly to the zip"""
+        # Create files
+        f1 = tmp_path / "file1.txt"
+        f1.write_text("file1 content")
+        f2 = tmp_path / "file2.txt"
+        f2.write_text("file2 content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), str(tmp_path), ["file1.txt", "file2.txt"])
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "file1.txt" in names
+            assert "file2.txt" in names
+            assert len(names) == 2
+
+    def test_directories_added_recursively(self, tmp_path):
+        """Ensure directories are added recursively with correct relative paths"""
+        nested = tmp_path / "folder"
+        nested.mkdir()
+        (nested / "nested1.txt").write_text("nested1 content")
+        subfolder = nested / "sub"
+        subfolder.mkdir()
+        (subfolder / "nested2.txt").write_text("nested2 content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), str(tmp_path), ["folder"])
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "folder/nested1.txt" in names
+            assert "folder/sub/nested2.txt" in names
+
+    def test_missing_files_skipped(self, tmp_path, caplog):
+        """Ensure missing files do not break the function and warning is logged"""
+        caplog.set_level(logging.WARNING)
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), str(tmp_path), ["does_not_exist.txt"], logger=logging.getLogger())
+
+        # Zip should exist but be empty
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            assert zf.namelist() == []
+
+        # Check warning was logged
+        assert any("Skipping missing file" in record.message for record in caplog.records)
+
+    def test_arcname_paths_correct(self, tmp_path):
+        """Ensure arcname paths inside zip preserve relative paths to base_path"""
+        (tmp_path / "file.txt").write_text("content")
+        folder = tmp_path / "folder"
+        folder.mkdir()
+        (folder / "nested.txt").write_text("nested content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), str(tmp_path), ["file.txt", "folder"])
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "file.txt" in names
+            assert "folder/nested.txt" in names

mito_ai/tests/providers/test_anthropic_client.py

@@ -53,7 +53,7 @@ def test_no_system_instructions_only_content():
     ]
     system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
     
-    assert isinstance(system_prompt, anthropic.NotGiven)
+    assert isinstance(system_prompt, anthropic.Omit)
     assert len(anthropic_messages) == 2
     assert anthropic_messages[0]["role"] == "user"
     assert anthropic_messages[0]["content"] == "Hello!"
@@ -93,7 +93,7 @@ def test_empty_message_content():
     ]
     system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
     
-    assert isinstance(system_prompt, anthropic.NotGiven)
+    assert isinstance(system_prompt, anthropic.Omit)
     assert len(anthropic_messages) == 1  # Should skip the message with missing content
     assert anthropic_messages[0]["role"] == "assistant"
     assert anthropic_messages[0]["content"] == "Hi!"

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -265,89 +265,5 @@ class TestStreamlitHandler:
         with pytest.raises(Exception, match="Generation failed"):
             await streamlit_handler("/path/to/notebook.ipynb")
 
-    @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.validate_app')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
-    async def test_streamlit_handler_too_many_files_in_directory(self, mock_clean_directory, mock_create_file, mock_validator, mock_generator_class, mock_parse):
-        """Test streamlit handler when there are too many files in the directory"""
-        # Mock clean directory check to raise ValueError (simulating >10 files)
-        mock_clean_directory.side_effect = ValueError("Too many files in directory: 10 allowed but 15 present. Create a new directory and retry")
-        
-        # The function should raise the ValueError before any other processing
-        with pytest.raises(ValueError, match="Too many files in directory: 10 allowed but 15 present. Create a new directory and retry"):
-            await streamlit_handler("/path/to/notebook.ipynb")
-        
-        # Verify that clean_directory_check was called
-        mock_clean_directory.assert_called_once_with("/path/to/notebook.ipynb")
-        
-        # Verify that no other functions were called since the error occurred early
-        mock_parse.assert_not_called()
-        mock_generator_class.assert_not_called()
-        mock_validator.assert_not_called()
-        mock_create_file.assert_not_called()
 
 
-class TestCleanDirectoryCheck:
-    """Test cases for clean_directory_check function"""
-
-    @patch('mito_ai.streamlit_conversion.streamlit_utils.Path')
-    def test_clean_directory_check_under_limit(self, mock_path):
-        """Test clean_directory_check when directory has 10 or fewer files"""
-        # Mock the Path class and its methods
-        mock_path_instance = mock_path.return_value
-        mock_path_instance.resolve.return_value = mock_path_instance
-        mock_path_instance.parent = mock_path_instance
-        
-        # Mock directory existence check
-        mock_path_instance.exists.return_value = True
-        
-        # Mock directory contents with 8 files
-        mock_files = []
-        for i in range(8):
-            mock_file = MagicMock()
-            mock_file.is_file.return_value = True
-            mock_files.append(mock_file)
-        
-        mock_path_instance.iterdir.return_value = mock_files
-        
-        # Should not raise any exception
-        clean_directory_check('/path/to/notebook.ipynb')
-        
-        # Verify calls
-        mock_path.assert_called_once_with('/path/to/notebook.ipynb')
-        mock_path_instance.resolve.assert_called_once()
-        mock_path_instance.exists.assert_called_once()
-        mock_path_instance.iterdir.assert_called_once()
-
-    @patch('mito_ai.streamlit_conversion.streamlit_utils.Path')
-    def test_clean_directory_check_over_limit(self, mock_path):
-        """Test clean_directory_check when directory has more than 10 files"""
-        # Mock the Path class and its methods
-        mock_path_instance = mock_path.return_value
-        mock_path_instance.resolve.return_value = mock_path_instance
-        mock_path_instance.parent = mock_path_instance
-        
-        # Mock directory existence check
-        mock_path_instance.exists.return_value = True
-        
-        # Mock directory contents with 15 files
-        mock_files = []
-        for i in range(15):
-            mock_file = MagicMock()
-            mock_file.is_file.return_value = True
-            mock_files.append(mock_file)
-        
-        mock_path_instance.iterdir.return_value = mock_files
-        
-        # Should raise ValueError
-        with pytest.raises(ValueError, match="Too many files in directory: 10 allowed but 15 present. Create a new directory and retry"):
-            clean_directory_check('/path/to/notebook.ipynb')
-        
-        # Verify calls
-        mock_path.assert_called_once_with('/path/to/notebook.ipynb')
-        mock_path_instance.resolve.assert_called_once()
-        mock_path_instance.exists.assert_called_once()
-        mock_path_instance.iterdir.assert_called_once()