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

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.44'
+__version__ = VERSION = '0.1.45'

mito_ai/anthropic_client.py

@@ -52,12 +52,12 @@ def extract_and_parse_anthropic_json_response(response: Message) -> Union[object
 
 
 def get_anthropic_system_prompt_and_messages(messages: List[ChatCompletionMessageParam]) -> Tuple[
-    Union[str, anthropic.NotGiven], List[MessageParam]]:
+    Union[str, anthropic.Omit], List[MessageParam]]:
     """
     Convert a list of OpenAI messages to a list of Anthropic messages.
     """
     
-    system_prompt: Union[str, anthropic.NotGiven] = anthropic.NotGiven()
+    system_prompt: Union[str, anthropic.Omit] = anthropic.Omit()
     anthropic_messages: List[MessageParam] = []
 
     for message in messages:
@@ -206,7 +206,6 @@ class AnthropicClient:
                     stream=True
                 )
 
-
                 for chunk in stream:
                     if chunk.type == "content_block_delta" and chunk.delta.type == "text_delta":
                         content = chunk.delta.text

mito_ai/app_deploy/app_deploy_utils.py

@@ -0,0 +1,25 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import zipfile
+import logging
+from typing import List, Optional
+
+def add_files_to_zip(zip_path: str, base_path: str, files_to_add: List[str], logger: Optional[logging.Logger] = None) -> None:
+    """Create a zip file at zip_path and add the selected files/folders."""
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for rel_path in files_to_add:
+            abs_path = os.path.join(base_path, rel_path)
+
+            if os.path.isfile(abs_path):
+                zipf.write(abs_path, arcname=rel_path)
+            elif os.path.isdir(abs_path):
+                for root, _, files in os.walk(abs_path):
+                    for file in files:
+                        file_abs = os.path.join(root, file)
+                        arcname = os.path.relpath(file_abs, base_path)
+                        zipf.write(file_abs, arcname=arcname)
+            else:
+                if logger:
+                    logger.warning(f"Skipping missing file: {abs_path}")

mito_ai/app_deploy/handlers.py

@@ -4,13 +4,13 @@
 import os
 import time
 import logging
-from typing import Any, Union, Optional
-import zipfile
+from typing import Any, Union, List
 import tempfile
 from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
 from mito_ai.utils.create import initialize_user
 from mito_ai.utils.version_utils import is_pro
 from mito_ai.utils.websocket_base import BaseWebSocketHandler
+from mito_ai.app_deploy.app_deploy_utils import  add_files_to_zip
 from mito_ai.app_deploy.models import (
     DeployAppReply,
     AppDeployError,
@@ -18,7 +18,6 @@ from mito_ai.app_deploy.models import (
     ErrorMessage,
     MessageType
 )
-from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
 from mito_ai.logger import get_logger
 from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
 import requests
@@ -111,6 +110,7 @@ class AppDeployHandler(BaseWebSocketHandler):
         message_id = message.message_id
         notebook_path = message.notebook_path
         jwt_token = message.jwt_token
+        files_to_upload = message.selected_files
         
         if not message_id:
             self.log.error("Missing message_id in request")
@@ -168,7 +168,7 @@ class AppDeployHandler(BaseWebSocketHandler):
                 ))
             
             # Finally, deploy the app
-            deploy_url = await self._deploy_app(app_directory, jwt_token)
+            deploy_url = await self._deploy_app(app_directory, files_to_upload, jwt_token)
 
             # Send the response
             self.reply(DeployAppReply(
@@ -219,11 +219,12 @@ class AppDeployHandler(BaseWebSocketHandler):
             return False
 
 
-    async def _deploy_app(self, app_path: str, jwt_token: str = '') -> str:
+    async def _deploy_app(self, app_path: str, files_to_upload:List[str], jwt_token: str = '') -> str:
         """Deploy the app using pre-signed URLs.
         
         Args:
             app_path: Path to the app file.
+            files_to_upload: Files the user selected to upload for the app to run
             jwt_token: JWT token for authentication (optional)
 
         Returns:
@@ -258,16 +259,12 @@ class AppDeployHandler(BaseWebSocketHandler):
             # Step 2: Create a zip file of the app.
             temp_zip_path = None
             try:
-                # Create temp file and close it before writing to avoid file handle conflicts
-                with tempfile.NamedTemporaryFile(suffix='.zip', delete=False) as temp_zip:
+                # Create temp file
+                with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as temp_zip:
                     temp_zip_path = temp_zip.name
 
                 self.log.info("Zipping application files...")
-                with zipfile.ZipFile(temp_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
-                    for root, _, files in os.walk(app_path):
-                        for file in files:
-                            file_path = os.path.join(root, file)
-                            zipf.write(file_path, arcname=os.path.relpath(file_path, app_path))
+                add_files_to_zip(temp_zip_path, app_path, files_to_upload, self.log)
 
                 upload_response = await self._upload_app_to_s3(temp_zip_path, presigned_url)
             except Exception as e:

mito_ai/app_deploy/models.py

@@ -3,7 +3,7 @@
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import Literal, Optional
+from typing import Literal, Optional, List
 
 
 class MessageType(str, Enum):
@@ -66,6 +66,9 @@ class DeployAppRequest:
     
     # Path to the app file.
     notebook_path: str
+
+    # Files to be uploaded for the app to run
+    selected_files: List[str]
     
     # JWT token for authorization.
     jwt_token: Optional[str] = None

mito_ai/completions/handlers.py

@@ -14,6 +14,7 @@ import tornado.web
 from jupyter_core.utils import ensure_async
 from jupyter_server.base.handlers import JupyterHandler
 from tornado.websocket import WebSocketHandler
+from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.logger import get_logger
 from mito_ai.completions.models import (
@@ -222,7 +223,32 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                 )
                 self.reply(reply)
             return
-        
+
+        if type == MessageType.STOP_AGENT:
+            thread_id_to_stop = metadata_dict.get('threadId')
+            if thread_id_to_stop:
+                self.log.info(f"Stopping agent, thread ID: {thread_id_to_stop}")
+                
+                ai_optimized_message: ChatCompletionMessageParam = {
+                    "role": "assistant",
+                    "content": "The user made the following request: Stop processing my last request. I want to change it. Please answer my future requests without going back and finising my previous request."
+                }
+                display_optimized_message: ChatCompletionMessageParam = {
+                    "role": "assistant",
+                    "content": "Agent interupted by user "
+                }
+                
+                await message_history.append_message(
+                    ai_optimized_message=ai_optimized_message,
+                    display_message=display_optimized_message,
+                    model=self._selected_model,
+                    llm_provider=self._llm,
+                    thread_id=thread_id_to_stop
+                )
+            else:
+                self.log.info("Trying to stop agent, but no thread ID available")
+            return
+
         try:
             # Get completion based on message type
             completion = None

mito_ai/completions/models.py

@@ -64,6 +64,7 @@ class MessageType(Enum):
     DELETE_THREAD = "delete_thread"
     UPDATE_MODEL_CONFIG = "update_model_config"
     STREAMLIT_CONVERSION = "streamlit_conversion"
+    STOP_AGENT = "stop_agent"
 
     
 @dataclass(frozen=True)

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -125,15 +125,33 @@ If the user has requested data that you believe is stored in the database:
 connections[connection_name]["username"]
 ```
 
+- The user may colloquially ask for a "list of x", always assume they want a pandas DataFrame. 
+- When working with dataframes created from an SQL query, ALWAYS use lowercase column names. 
+- If you think the requested data is stored in the database, but you are unsure, then ask the user for clarification.
+
+## Additional MSSQL Rules
+
+- When connecting to a Microsoft SQL Server (MSSQL) database, use the following format:
+
+```
+import urllib.parse
+
+encoded_password = urllib.parse.quote_plus(password) 
+conn_str = f"mssql+pyodbc://username:encoded_password@host:port/database?driver=ODBC+Driver+18+for+SQL+Server&TrustServerCertificate=yes"
+```
+
+- Always URL-encode passwords for MSSQL connections to handle special characters properly.
+- Include the port number in MSSQL connection strings.
+- Use "ODBC+Driver+18+for+SQL+Server" (with plus signs) in the driver parameter.
+- Always include "TrustServerCertificate=yes" for MSSQL connections to avoid SSL certificate issues.
+
+## Additional Oracle Rules
+
 - When connecting to an Oracle database, use the following format:
 ```
 conn_str = f"oracle+oracledb://username:password@host:port?service_name=service_name"
 ```
 
-- The user may colloquially ask for a "list of x", always assume they want a pandas DataFrame. 
-- When working with dataframes created from an SQL query, ALWAYS use lowercase column names. 
-- If you think the requested data is stored in the database, but you are unsure, then ask the user for clarification.
-
 Here is the schema:
 {schemas}
         """

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -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,8 +54,10 @@ 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)}"
@@ -139,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

@@ -88,9 +88,12 @@ class StreamlitPreviewHandler(APIHandler):
                 self.finish({"error": error_msg})
                 return
 
+
             # Ensure app exists
             resolved_notebook_path = self._resolve_notebook_path(notebook_path)
+
             success, error_msg = await ensure_app_exists(resolved_notebook_path, force_recreate)
+
             if not success:
                 self.set_status(500)
                 self.finish({"error": error_msg})

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()

mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py

@@ -87,21 +87,6 @@ df=pd.read_csv('data.csv')
             validator = StreamlitValidator() 
             errors = validator.get_runtime_errors(app_code, app_path)
             assert errors is None
-            
-            
-    @patch('subprocess.Popen')
-    def test_cleanup_with_process(self, mock_popen):
-        """Test cleanup with running process"""
-        validator = StreamlitValidator()
-        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_rmtree.assert_called_once()
-
 
     @pytest.mark.parametrize("app_code,expected_has_validation_error,expected_error_message", [
         ("x=5", False, ""),

mito_ai/tests/utils/test_anthropic_utils.py

@@ -35,7 +35,7 @@ def test_basic_request_preparation():
     max_tokens = 100
     temperature = 0.7
     # Use NotGiven to ensure system is not included in inner_data
-    system = anthropic.NotGiven()
+    system = anthropic.Omit()
     messages: List[MessageParam] = [{"role": "user", "content": "Hello"}]
     message_type = MessageType.CHAT
 
@@ -106,7 +106,7 @@ def test_tools_and_tool_choice():
         model="claude-3-sonnet",
         max_tokens=100,
         temperature=0.7,
-        system=anthropic.NotGiven(),
+        system=anthropic.Omit(),
         messages=[{"role": "user", "content": "Hello"}],
         message_type=MessageType.CHAT,
         tools=tools,
@@ -124,7 +124,7 @@ def test_stream_parameter():
         model="claude-3-sonnet",
         max_tokens=100,
         temperature=0.7,
-        system=anthropic.NotGiven(),
+        system=anthropic.Omit(),
         messages=[{"role": "user", "content": "Hello"}],
         message_type=MessageType.CHAT,
         tools=None,
@@ -150,7 +150,7 @@ def test_missing_user_info(monkeypatch):
         model="claude-3-sonnet",
         max_tokens=100,
         temperature=0.7,
-        system=anthropic.NotGiven(),
+        system=anthropic.Omit(),
         messages=[{"role": "user", "content": "Hello"}],
         message_type=MessageType.CHAT,
         tools=None,