mito-ai 0.1.47__py3-none-any.whl → 0.1.48__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.47'
+__version__ = VERSION = '0.1.48'

mito_ai/app_deploy/handlers.py

@@ -6,7 +6,7 @@ import time
 import logging
 from typing import Any, Union, List, Optional
 import tempfile
-from mito_ai.path_utils import AbsoluteAppPath, does_app_path_exists, get_absolute_app_path, get_absolute_notebook_dir_path, get_absolute_notebook_path
+from mito_ai.path_utils import AbsoluteAppPath, does_app_path_exist, get_absolute_app_path, get_absolute_notebook_dir_path, get_absolute_notebook_path
 from mito_ai.utils.create import initialize_user
 from mito_ai.utils.error_classes import StreamlitDeploymentError
 from mito_ai.utils.version_utils import is_pro
@@ -131,6 +131,7 @@ class AppDeployHandler(BaseWebSocketHandler):
             raise StreamlitDeploymentError(error)
 
         if not notebook_path:
+            self.log.error("Missing notebook_path in request")
             error = AppDeployError(
                 error_type="InvalidRequest",
                 message="Missing 'notebook_path' parameter",
@@ -162,7 +163,7 @@ class AppDeployHandler(BaseWebSocketHandler):
         app_path = get_absolute_app_path(absolute_app_directory)
 
         # Check if the app.py file exists
-        app_path_exists = does_app_path_exists(app_path)
+        app_path_exists = does_app_path_exist(app_path)
         if not app_path_exists:
             error = AppDeployError(
                 error_type="AppNotFound",

mito_ai/completions/models.py

@@ -97,6 +97,7 @@ class AgentExecutionMetadata():
     threadId: ThreadID
     input: str
     aiOptimizedCells: List[AIOptimizedCell]
+    activeCellId: str
     isChromeBrowser: bool
     base64EncodedActiveCellOutput: Optional[str] = None
     variables: Optional[List[str]] = None

mito_ai/completions/prompt_builders/agent_execution_prompt.py

@@ -3,6 +3,7 @@
 
 from mito_ai.completions.models import AgentExecutionMetadata
 from mito_ai.completions.prompt_builders.prompt_constants import (
+    ACTIVE_CELL_ID_SECTION_HEADING,
     FILES_SECTION_HEADING,
     JUPYTER_NOTEBOOK_SECTION_HEADING,
     STREAMLIT_APP_STATUS_SECTION_HEADING,
@@ -40,6 +41,9 @@ def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
 {STREAMLIT_APP_STATUS_SECTION_HEADING}
 {streamlit_status_str}
 
+{ACTIVE_CELL_ID_SECTION_HEADING}
+{md.activeCellId}
+
 {selected_context_str}
 
 {cell_update_output_str(md.base64EncodedActiveCellOutput is not None)}"""

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -469,4 +469,4 @@ REMEMBER, YOU ARE GOING TO COMPLETE THE USER'S TASK OVER THE COURSE OF THE ENTIR
 
 OTHER USEFUL INFORMATION:
 1. When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter
-"""
+2. The active cell ID is shared with you so that when the user refers to "this cell" or similar phrases, you know which cell they mean. However, you are free to edit any cell that you see fit."""

mito_ai/path_utils.py

@@ -48,7 +48,7 @@ def get_absolute_app_path(app_directory: AbsoluteNotebookDirPath) -> AbsoluteApp
     """
     return AbsoluteAppPath(os.path.join(app_directory, "app.py"))
 
-def does_app_path_exists(app_path: AbsoluteAppPath) -> bool:
+def does_app_path_exist(app_path: AbsoluteAppPath) -> bool:
     """
     Check if the app.py file exists in the given directory.
     """

mito_ai/streamlit_conversion/search_replace_utils.py

@@ -5,6 +5,7 @@ import re
 from typing import List, Tuple
 
 from mito_ai.utils.error_classes import StreamlitConversionError
+from mito_ai.utils.telemetry_utils import log
 
 
 def extract_search_replace_blocks(message_content: str) -> List[Tuple[str, str]]:
@@ -85,9 +86,9 @@ def apply_search_replace(text: str, search_replace_pairs: List[Tuple[str, str]])
             raise StreamlitConversionError(f"Search text not found: {repr(search_text)}", error_code=500)
         elif count > 1:
             print("Search Text Found Multiple Times: ", repr(search_text))
-            raise StreamlitConversionError(f"Search text found {count} times (must be unique): {repr(search_text)}", error_code=500)
-        
+            log("mito_ai_search_text_found_multiple_times", params={"search_text": repr(search_text)}, key_type="mito_server_key")
+
         # Perform the replacement
-        result = result.replace(search_text, replace_text)
+        result = result.replace(search_text, replace_text, 1)
     
     return result

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -1,9 +1,8 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-import os
 from anthropic.types import MessageParam
-from typing import List, Optional, Tuple, cast
+from typing import List, cast
 from mito_ai.streamlit_conversion.agent_utils import extract_todo_placeholders, get_response_from_agent
 from mito_ai.streamlit_conversion.prompts.streamlit_app_creation_prompt import get_streamlit_app_creation_prompt
 from mito_ai.streamlit_conversion.prompts.streamlit_error_correction_prompt import get_streamlit_error_correction_prompt
@@ -122,21 +121,21 @@ async def streamlit_handler(notebook_path: AbsoluteNotebookPath, edit_prompt: st
         streamlit_code = await generate_new_streamlit_code(notebook_code)
        
     # Then, after creating/updating the app, validate that the new code runs 
-    has_validation_error, errors = validate_app(streamlit_code, notebook_path)
+    errors = validate_app(streamlit_code, notebook_path)
     tries = 0
-    while has_validation_error and tries < 5:
+    while len(errors)>0 and tries < 5:
         for error in errors:
             streamlit_code = await correct_error_in_generation(error, streamlit_code)
         
-        has_validation_error, errors = validate_app(streamlit_code, notebook_path)
+        errors = validate_app(streamlit_code, notebook_path)
         
-        if has_validation_error:
+        if len(errors)>0:
             # TODO: We can't easily get the key type here, so for the beta release
             # we are just defaulting to the mito server key since that is by far the most common.
             log_streamlit_app_validation_retry('mito_server_key', MessageType.STREAMLIT_CONVERSION, errors)
         tries+=1
 
-    if has_validation_error:
+    if len(errors)>0:
         final_errors = ', '.join(errors)
         raise StreamlitConversionError(f"Streamlit agent failed generating code after max retries. Errors: {final_errors}", 500)
     

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -3,9 +3,8 @@
 
 import re
 import json
-from typing import Dict, List, Optional, Tuple, Any
+from typing import Dict, List, Optional, Any
 from mito_ai.path_utils import AbsoluteAppPath, AbsoluteNotebookPath
-from pathlib import Path
 from mito_ai.utils.error_classes import StreamlitConversionError
 
 

mito_ai/streamlit_conversion/validate_streamlit_app.py

@@ -98,10 +98,8 @@ def check_for_errors(app_code: str, app_path: AbsoluteNotebookPath) -> List[Dict
 
     return errors
 
-def validate_app(app_code: str, notebook_path: AbsoluteNotebookPath) -> Tuple[bool, List[str]]:
+def validate_app(app_code: str, notebook_path: AbsoluteNotebookPath) -> List[str]:
     """Convenience function to validate Streamlit code"""
     errors = check_for_errors(app_code, notebook_path)
-    
-    has_validation_error = len(errors) > 0
     stringified_errors = [str(error) for error in errors]
-    return has_validation_error, stringified_errors
+    return stringified_errors

mito_ai/streamlit_preview/__init__.py

@@ -1,7 +1,6 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from .manager import get_preview_manager
 from .handlers import StreamlitPreviewHandler
 
-__all__ = ['get_preview_manager', 'StreamlitPreviewHandler'] 
+__all__ = ['StreamlitPreviewHandler']

mito_ai/streamlit_preview/handlers.py

@@ -1,17 +1,16 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-import os
-from typing import Literal, TypedDict
 import uuid
-from mito_ai.streamlit_preview.utils import ensure_app_exists, validate_request_body
+from mito_ai.streamlit_preview.utils import validate_request_body
 import tornado
 from jupyter_server.base.handlers import APIHandler
-from mito_ai.streamlit_preview.manager import get_preview_manager
-from mito_ai.path_utils import  AbsoluteNotebookPath, get_absolute_notebook_dir_path, get_absolute_notebook_path
+from mito_ai.streamlit_preview.manager import StreamlitPreviewManager
+from mito_ai.path_utils import get_absolute_notebook_dir_path, get_absolute_notebook_path, get_absolute_app_path, does_app_path_exist
 from mito_ai.utils.telemetry_utils import log_streamlit_app_conversion_error, log_streamlit_app_preview_failure, log_streamlit_app_preview_success
 from mito_ai.completions.models import MessageType
 from mito_ai.utils.error_classes import StreamlitConversionError, StreamlitPreviewError
+from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
 import traceback
 
 
@@ -20,7 +19,7 @@ class StreamlitPreviewHandler(APIHandler):
 
     def initialize(self) -> None:
         """Initialize the handler."""
-        self.preview_manager = get_preview_manager()
+        self.preview_manager = StreamlitPreviewManager()
 
     @tornado.web.authenticated
     async def post(self) -> None:
@@ -32,15 +31,24 @@ class StreamlitPreviewHandler(APIHandler):
 
             # Ensure app exists
             absolute_notebook_path = get_absolute_notebook_path(notebook_path)
-            await ensure_app_exists(absolute_notebook_path, force_recreate, edit_prompt)
+            absolute_notebook_dir_path = get_absolute_notebook_dir_path(absolute_notebook_path)
+            absolute_app_path = get_absolute_app_path(absolute_notebook_dir_path)
+            app_path_exists = does_app_path_exist(absolute_app_path)
+
+            if not app_path_exists or force_recreate:
+                if not app_path_exists:
+                    print("[Mito AI] App path not found, generating streamlit code")
+                else:
+                    print("[Mito AI] Force recreating streamlit app")
+
+                await streamlit_handler(absolute_notebook_path, edit_prompt)
 
             # 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())
-            absolute_app_directory = get_absolute_notebook_dir_path(absolute_notebook_path)
-            port = self.preview_manager.start_streamlit_preview(absolute_app_directory, preview_id)
+            port = self.preview_manager.start_streamlit_preview(absolute_notebook_dir_path, preview_id)
 
             # Return success response
             self.finish({

mito_ai/streamlit_preview/manager.py

@@ -1,10 +1,8 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-import os
 import socket
 import subprocess
-import tempfile
 import time
 import threading
 import requests
@@ -107,7 +105,7 @@ class StreamlitPreviewManager:
                 if response.status_code == 200:
                     return True
             except requests.exceptions.RequestException as e:
-                print(f"Waiting for app to be ready...")
+                self.log.info(f"Waiting for app to be ready...")
                 pass
             
             time.sleep(1)
@@ -123,7 +121,7 @@ class StreamlitPreviewManager:
         Returns:
             True if stopped successfully, False if not found
         """
-        print(f"Stopping preview {preview_id}")
+        self.log.info(f"Stopping preview {preview_id}")
         with self._lock:
             if preview_id not in self._previews:
                 return False
@@ -150,10 +148,3 @@ class StreamlitPreviewManager:
         """Get a preview process by ID."""
         with self._lock:
             return self._previews.get(preview_id)
-
-# Global instance
-_preview_manager = StreamlitPreviewManager()
-
-def get_preview_manager() -> StreamlitPreviewManager:
-    """Get the global preview manager instance."""
-    return _preview_manager 

mito_ai/streamlit_preview/utils.py

@@ -2,8 +2,6 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 from typing import Tuple, Optional
-from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
-from mito_ai.path_utils import AbsoluteNotebookPath, does_app_path_exists, get_absolute_app_path, get_absolute_notebook_dir_path
 from mito_ai.utils.error_classes import StreamlitPreviewError
 
 
@@ -19,24 +17,9 @@ def validate_request_body(body: Optional[dict]) -> Tuple[str, bool, str]:
     force_recreate = body.get("force_recreate", False)
     if not isinstance(force_recreate, bool):
         raise StreamlitPreviewError("force_recreate must be a boolean", 400)
-    
+
     edit_prompt = body.get("edit_prompt", "")
     if not isinstance(edit_prompt, str):
         raise StreamlitPreviewError("edit_prompt must be a string", 400)
 
     return notebook_path, force_recreate, edit_prompt
-
-async def ensure_app_exists(absolute_notebook_path: AbsoluteNotebookPath, force_recreate: bool = False, edit_prompt: str = "") -> None:
-    """Ensure app.py exists, generating it if necessary or if force_recreate is True."""
-
-    absolute_notebook_dir_path = get_absolute_notebook_dir_path(absolute_notebook_path)
-    absolute_app_path = get_absolute_app_path(absolute_notebook_dir_path)
-    app_path_exists = does_app_path_exists(absolute_app_path)
-    
-    if not app_path_exists or force_recreate:
-        if not app_path_exists:
-            print("[Mito AI] App path not found, generating streamlit code")
-        else:
-            print("[Mito AI] Force recreating streamlit app")
-        
-        await streamlit_handler(absolute_notebook_path, edit_prompt)

mito_ai/tests/message_history/test_message_history_utils.py

@@ -112,6 +112,7 @@ PROMPT_BUILDER_TEST_CASES = [
                 input=TEST_INPUT,
                 promptType="agent:execution",
                 threadId=ThreadID("test-thread-id"),
+                activeCellId="cell1",
                 isChromeBrowser=True
             )
         ),

mito_ai/tests/streamlit_conversion/test_apply_search_replace.py

@@ -200,6 +200,23 @@ st.write("Welcome to my application")""",
 
 st.title("🚀 My App")
 st.write("Welcome to my application")"""
+    ),
+    
+    # Test case 11: Only replace first occurrence when search text exists multiple times
+    (
+        """import streamlit as st
+
+st.write("Hello World")
+st.title("My App")
+st.write("Hello World")
+st.write("Another message")""",
+        [("st.write(\"Hello World\")", "st.write(\"Hi There\")")],
+        """import streamlit as st
+
+st.write("Hi There")
+st.title("My App")
+st.write("Hello World")
+st.write("Another message")"""
     )
 ])
 def test_apply_search_replace(original_text, search_replace_pairs, expected_result):
@@ -220,7 +237,4 @@ def test_apply_search_replace_search_not_found():
         apply_search_replace("st.title(\"My App\")", [("st.title(\"Not Found\")", "st.title(\"New Title\")")])
 
 
-def test_apply_search_replace_multiple_matches():
-    """Test that ValueError is raised when search text is found multiple times."""
-    with pytest.raises(StreamlitConversionError, match="Search text found 2 times"):
-        apply_search_replace("st.write(\"Hello\")\nst.write(\"Hello\")", [("st.write(\"Hello\")", "st.write(\"Hi\")")])
+

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -138,7 +138,8 @@ class TestStreamlitHandler:
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
-    async def test_streamlit_handler_success(self, mock_create_file, mock_validator, mock_generate_code, mock_parse):
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.log_streamlit_app_conversion_success')
+    async def test_streamlit_handler_success(self, mock_log_success, mock_create_file, mock_validator, mock_generate_code, mock_parse):
         """Test successful streamlit handler execution"""
         # Mock notebook parsing
         mock_notebook_data: List[dict] = [{"cells": [{"cell_type": "code", "source": ["import pandas"]}]}]
@@ -148,7 +149,7 @@ class TestStreamlitHandler:
         mock_generate_code.return_value = "import streamlit\nst.title('Test')"
         
         # Mock validation (no errors)
-        mock_validator.return_value = (False, "")
+        mock_validator.return_value = []
         
         # Use a relative path that will work cross-platform
         notebook_path = AbsoluteNotebookPath("absolute/path/to/notebook.ipynb")
@@ -169,7 +170,8 @@ class TestStreamlitHandler:
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.correct_error_in_generation')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
-    async def test_streamlit_handler_max_retries_exceeded(self, mock_validator, mock_correct_error, mock_generate_code, mock_parse):
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.log_streamlit_app_validation_retry')
+    async def test_streamlit_handler_max_retries_exceeded(self, mock_log_retry, mock_validator, mock_correct_error, mock_generate_code, mock_parse):
         """Test streamlit handler when max retries are exceeded"""
         # Mock notebook parsing
         mock_notebook_data: List[dict] = [{"cells": []}]
@@ -179,14 +181,15 @@ class TestStreamlitHandler:
         mock_generate_code.return_value = "import streamlit\nst.title('Test')"
         mock_correct_error.return_value = "import streamlit\nst.title('Fixed')"
     
-        # Mock validation (always errors) - Return list of errors as expected by validate_app
-        mock_validator.return_value = (True, ["Persistent error"])
+        # Mock validation (always errors) - validate_app returns List[str]
+        mock_validator.return_value = ["Persistent error"]
     
         # Now it should raise an exception instead of returning a tuple
         with pytest.raises(Exception):
             await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"))
         
-        # Verify that error correction was called 5 times (max retries)
+        # Verify that error correction was called 5 times (once per error, 5 retries)
+        # Each retry processes 1 error, so 5 retries = 5 calls
         assert mock_correct_error.call_count == 5
 
     @pytest.mark.asyncio
@@ -203,14 +206,14 @@ class TestStreamlitHandler:
         # Mock code generation
         mock_generate_code.return_value = "import streamlit\nst.title('Test')"
         
-        # Mock validation (no errors)
-        mock_validator.return_value = (False, "")
+        # Mock validation (no errors) - validate_app returns List[str]
+        mock_validator.return_value = []
         
         # Mock file creation failure - now it should raise an exception
         mock_create_file.side_effect = Exception("Permission denied")
         
         # Now it should raise an exception instead of returning a tuple
-        with pytest.raises(Exception, match="Permission denied"):
+        with pytest.raises(Exception):
             await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"))
 
     @pytest.mark.asyncio

mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py

@@ -7,7 +7,6 @@ from unittest.mock import patch, MagicMock
 from mito_ai.streamlit_conversion.validate_streamlit_app import (
     get_syntax_error,
     get_runtime_errors,
-    check_for_errors,
     validate_app
 )
 import pytest
@@ -20,33 +19,33 @@ class TestGetSyntaxError:
     @pytest.mark.parametrize("code,expected_error,test_description", [
         # Valid Python code should return no error
         (
-            "import streamlit\nst.title('Hello World')",
-            None,
-            "valid Python code"
+                "import streamlit\nst.title('Hello World')",
+                None,
+                "valid Python code"
         ),
         # Invalid Python syntax should be caught
         (
-            "import streamlit\nst.title('Hello World'",
-            "SyntaxError",
-            "invalid Python code"
+                "import streamlit\nst.title('Hello World'",
+                "SyntaxError",
+                "invalid Python code"
         ),
         # Empty streamlit app is valid
         (
-            "",
-            None,
-            "empty code"
+                "",
+                None,
+                "empty code"
         ),
     ])
     def test_get_syntax_error(self, code, expected_error, test_description):
         """Test syntax validation with various code inputs"""
         error = get_syntax_error(code)
-        
+
         if expected_error is None:
             assert error is None, f"Expected no error for {test_description}"
         else:
             assert error is not None, f"Expected error for {test_description}"
             assert expected_error in error, f"Expected '{expected_error}' in error for {test_description}"
-        
+
 class TestGetRuntimeErrors:
     """Test cases for get_runtime_errors function"""
 
@@ -94,17 +93,20 @@ df=pd.read_csv('data.csv')
 class TestValidateApp:
     """Test cases for validate_app function"""
 
-    @pytest.mark.parametrize("app_code,expected_has_validation_error,expected_error_message", [
+    @pytest.mark.parametrize("app_code,expected_has_errors,expected_error_message", [
         ("x=5", False, ""),
         ("1/0", True, "division by zero"),
-        ("print('Hello World'", True, "SyntaxError"),
+        # Syntax errors are caught during runtime
         ("", False, ""),
     ])
-    def test_validate_app(self, app_code, expected_has_validation_error, expected_error_message):
+    def test_validate_app(self, app_code, expected_has_errors, expected_error_message):
         """Test the validate_app function"""
-        has_validation_error, errors = validate_app(app_code, AbsoluteNotebookPath('/notebook.ipynb'))
+        errors = validate_app(app_code, AbsoluteNotebookPath('/notebook.ipynb'))
         
-        assert has_validation_error == expected_has_validation_error
-        assert expected_error_message in str(errors)
+        has_errors = len(errors) > 0
+        assert has_errors == expected_has_errors
+        if expected_error_message:
+            errors_str = str(errors)
+            assert expected_error_message in errors_str