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

mito_ai/streamlit_conversion/validate_streamlit_app.py

@@ -23,102 +23,97 @@ from mito_ai.streamlit_conversion.streamlit_utils import resolve_notebook_path
 warnings.filterwarnings("ignore", message=".*bare mode.*")
 
 
-class StreamlitValidator:
-    def __init__(self, port: int = 8501) -> None:
-        pass
+def get_syntax_error(app_code: str) -> Optional[str]:
+    """Check if the Python code has valid syntax"""
+    try:
+        ast.parse(app_code)
+        return None
+    except SyntaxError as e:
+        error_msg = ''.join(traceback.format_exception(type(e), e, e.__traceback__))
+        return error_msg
 
-    def get_syntax_error(self, app_code: str) -> Optional[str]:
-        """Check if the Python code has valid syntax"""
-        try:
-            ast.parse(app_code)
-            return None
-        except SyntaxError as e:
-            error_msg = ''.join(traceback.format_exception(type(e), e, e.__traceback__))
-            return error_msg
-
-    def get_runtime_errors(self, app_code: str, app_path: str) -> Optional[List[Dict[str, Any]]]:
-        """Start the Streamlit app in a subprocess"""  
-        
-        directory = os.path.dirname(app_path)
-        
-        @contextmanager
-        def change_working_directory(path: str) -> Generator[None, Any, None]:
-            """
-            Context manager to temporarily change working directory
-            so that relative paths are still valid when we run the app
-            """
-            if path == '':
-                yield
-            
-            original_cwd = os.getcwd()
-            try:
-                os.chdir(path)
-                yield
-            finally:
-                os.chdir(original_cwd)
+def get_runtime_errors(app_code: str, app_path: str) -> Optional[List[Dict[str, Any]]]:
+    """Start the Streamlit app in a subprocess"""  
+    
+    directory = os.path.dirname(app_path)
+    
+    @contextmanager
+    def change_working_directory(path: str) -> Generator[None, Any, None]:
+        """
+        Context manager to temporarily change working directory
+        so that relative paths are still valid when we run the app
+        """
+        if path == '':
+            yield
         
-        with change_working_directory(directory):
-            # 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. 
+        original_cwd = os.getcwd()
+        try:
+            os.chdir(path)
+            yield
+        finally:
+            os.chdir(original_cwd)
+    
+    with change_working_directory(directory):
+        # 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
+        # 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
 
+        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:
-                # 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
+                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"""
-        errors: List[Dict[str, Any]] = []
+def check_for_errors(app_code: str, app_path: str) -> List[Dict[str, Any]]:
+    """Complete validation pipeline"""
+    errors: List[Dict[str, Any]] = []
 
-        try:
-            # Step 1: Check syntax
-            syntax_error = self.get_syntax_error(app_code)
-            if syntax_error:
-                errors.append({'type': 'syntax', 'details': syntax_error})
+    try:
+        # Step 1: Check syntax
+        syntax_error = get_syntax_error(app_code)
+        if syntax_error:
+            errors.append({'type': 'syntax', 'details': syntax_error})
 
-            runtime_errors = self.get_runtime_errors(app_code, app_path)
-            
-            if runtime_errors:
-                errors.extend(runtime_errors)
-            
-        except Exception as e:
-            errors.append({'type': 'validation', 'details': str(e)})
+        runtime_errors = get_runtime_errors(app_code, app_path)
+        
+        if runtime_errors:
+            errors.extend(runtime_errors)
+        
+    except Exception as e:
+        errors.append({'type': 'validation', 'details': str(e)})
 
-        return errors
+    return errors
 
 def validate_app(app_code: str, notebook_path: str) -> Tuple[bool, List[str]]:
     """Convenience function to validate Streamlit code"""
     notebook_path = resolve_notebook_path(notebook_path)
     
-    validator = StreamlitValidator()
-    errors = validator._validate_app(app_code, notebook_path)
+    errors = check_for_errors(app_code, notebook_path)
     
     has_validation_error = len(errors) > 0
     stringified_errors = [str(error) for error in errors]

mito_ai/streamlit_preview/handlers.py

@@ -82,17 +82,16 @@ class StreamlitPreviewHandler(APIHandler):
         try:
             # Parse and validate request
             body = self.get_json_body()
-            is_valid, error_msg, notebook_path, force_recreate = validate_request_body(body)
-            if not is_valid or notebook_path is None:
+            is_valid, error_msg, notebook_path, force_recreate, edit_prompt = validate_request_body(body)
+            if not is_valid or not notebook_path:
                 self.set_status(400)
                 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)
+            success, error_msg = await ensure_app_exists(resolved_notebook_path, force_recreate, edit_prompt)
 
             if not success:
                 self.set_status(500)

mito_ai/streamlit_preview/utils.py

@@ -7,22 +7,26 @@ 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]:
+def validate_request_body(body: Optional[dict]) -> Tuple[bool, str, Optional[str], bool, str]:
     """Validate the request body and extract notebook_path and force_recreate."""
     if body is None:
-        return False, "Invalid or missing JSON body", None, False
+        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
+        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 False, "force_recreate must be a boolean", None, False, ""
+    
+    edit_prompt = body.get("edit_prompt", "")
+    if not isinstance(edit_prompt, str):
+        return False, "edit_prompt must be a string", None, False, ""
 
-    return True, "", notebook_path, force_recreate
+    return True, "", notebook_path, force_recreate, edit_prompt
 
-async def ensure_app_exists(resolved_notebook_path: str, force_recreate: bool = False) -> Tuple[bool, str]:
+async def ensure_app_exists(resolved_notebook_path: str, force_recreate: bool = False, edit_prompt: str = "") -> 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))
@@ -33,7 +37,7 @@ async def ensure_app_exists(resolved_notebook_path: str, force_recreate: bool =
         else:
             print("[Mito AI] Force recreating streamlit app")
         
-        success, app_path, message = await streamlit_handler(resolved_notebook_path)
+        success, app_path, message = await streamlit_handler(resolved_notebook_path, edit_prompt)
 
         if not success or app_path is None:
             return False, f"Failed to generate streamlit code: {message}"

mito_ai/tests/chat_history/test_chat_history.py

@@ -0,0 +1,211 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import requests
+import time
+from unittest.mock import patch, MagicMock
+from mito_ai.tests.conftest import TOKEN
+from mito_ai.completions.message_history import GlobalMessageHistory, ChatThread
+from mito_ai.completions.models import ThreadID
+
+
+@pytest.fixture
+def mock_chat_threads():
+    """Fixture that creates mock chat threads for testing"""
+    thread_id_1 = ThreadID("test-thread-1")
+    thread_id_2 = ThreadID("test-thread-2")
+
+    # Create mock threads with different timestamps
+    thread_1 = ChatThread(
+        thread_id=thread_id_1,
+        creation_ts=time.time() - 3600,  # 1 hour ago
+        last_interaction_ts=time.time() - 1800,  # 30 minutes ago
+        name="Test Chat 1",
+        ai_optimized_history=[
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there!"},
+        ],
+        display_history=[
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there!"},
+        ],
+    )
+
+    thread_2 = ChatThread(
+        thread_id=thread_id_2,
+        creation_ts=time.time() - 7200,  # 2 hours ago
+        last_interaction_ts=time.time() - 900,  # 15 minutes ago (more recent)
+        name="Test Chat 2",
+        ai_optimized_history=[
+            {"role": "user", "content": "How are you?"},
+            {"role": "assistant", "content": "I'm doing well, thanks!"},
+        ],
+        display_history=[
+            {"role": "user", "content": "How are you?"},
+            {"role": "assistant", "content": "I'm doing well, thanks!"},
+        ],
+    )
+
+    return {thread_id_1: thread_1, thread_id_2: thread_2}
+
+
+@pytest.fixture
+def mock_message_history(mock_chat_threads):
+    """Fixture that mocks the GlobalMessageHistory with test data"""
+    mock_history = MagicMock(spec=GlobalMessageHistory)
+    mock_history._chat_threads = mock_chat_threads
+
+    # Mock the get_threads method to return sorted threads
+    def mock_get_threads():
+        from mito_ai.completions.models import ChatThreadMetadata
+
+        threads = []
+        for thread in mock_chat_threads.values():
+            threads.append(
+                ChatThreadMetadata(
+                    thread_id=thread.thread_id,
+                    name=thread.name,
+                    creation_ts=thread.creation_ts,
+                    last_interaction_ts=thread.last_interaction_ts,
+                )
+            )
+        # Sort by last_interaction_ts (newest first)
+        threads.sort(key=lambda x: x.last_interaction_ts, reverse=True)
+        return threads
+
+    mock_history.get_threads = mock_get_threads
+    return mock_history
+
+
+# --- GET ALL THREADS ---
+
+
+def test_get_all_threads_success(jp_base_url: str, mock_message_history):
+    """Test successful GET all threads endpoint"""
+    # Since the server extension is already loaded, we need to work with the actual instance
+    # Let's just test that the endpoint works and returns the expected structure
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert "threads" in response_json
+    # The actual number of threads will depend on what's in the .mito/ai-chats directory
+    # So we'll just check that it's a list
+    assert isinstance(response_json["threads"], list)
+
+    # Check thread structure for any threads that exist
+    for thread in response_json["threads"]:
+        assert "thread_id" in thread
+        assert "name" in thread
+        assert "creation_ts" in thread
+        assert "last_interaction_ts" in thread
+
+
+def test_get_all_threads_empty(jp_base_url: str):
+    """Test GET all threads endpoint when no threads exist"""
+    # This test will work with whatever threads exist in the actual .mito/ai-chats directory
+    # We'll just verify the endpoint works and returns the expected structure
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert "threads" in response_json
+    assert isinstance(response_json["threads"], list)
+
+
+def test_get_all_threads_with_no_auth(jp_base_url: str):
+    """Test GET all threads endpoint without authentication"""
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_all_threads_with_incorrect_auth(jp_base_url: str):
+    """Test GET all threads endpoint with incorrect authentication"""
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+# --- GET SPECIFIC THREAD ---
+
+
+def test_get_specific_thread_success(jp_base_url: str, mock_message_history):
+    """Test successful GET specific thread endpoint"""
+    # First, get all threads to see what's available
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    threads = response.json()["threads"]
+    if not threads:
+        # If no threads exist, skip this test
+        pytest.skip("No threads available for testing")
+
+    # Use the first available thread
+    thread_id = threads[0]["thread_id"]
+
+    response = requests.get(
+        jp_base_url + f"/mito-ai/chat-history/threads/{thread_id}",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert response_json["thread_id"] == thread_id
+    assert "name" in response_json
+    assert "creation_ts" in response_json
+    assert "last_interaction_ts" in response_json
+    assert "display_history" in response_json
+    assert "ai_optimized_history" in response_json
+
+    # Check message history structure
+    display_history = response_json["display_history"]
+    assert isinstance(display_history, list)
+    ai_optimized_history = response_json["ai_optimized_history"]
+    assert isinstance(ai_optimized_history, list)
+
+
+def test_get_specific_thread_not_found(jp_base_url: str, mock_message_history):
+    """Test GET specific thread endpoint with non-existent thread ID"""
+    # Use a clearly non-existent thread ID
+    fake_thread_id = "non-existent-thread-12345"
+
+    response = requests.get(
+        jp_base_url + f"/mito-ai/chat-history/threads/{fake_thread_id}",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 404
+
+    response_json = response.json()
+    assert "error" in response_json
+    assert fake_thread_id in response_json["error"]
+
+
+def test_get_specific_thread_with_no_auth(jp_base_url: str):
+    """Test GET specific thread endpoint without authentication"""
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads/test-thread-1",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_specific_thread_with_incorrect_auth(jp_base_url: str):
+    """Test GET specific thread endpoint with incorrect authentication"""
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads/test-thread-1",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden

mito_ai/tests/message_history/test_message_history_utils.py

@@ -233,16 +233,17 @@ file1.csv
 file2.txt
 """
     
-    # Create test messages with proper typing
+    # Create test messages with proper typing 
     messages: List[ChatCompletionMessageParam] = [
         {"role": "system", "content": system_message_with_sections},
         {"role": "user", "content": user_message_with_sections},
         {"role": "assistant", "content": assistant_message_with_sections},
-        {"role": "user", "content": "Recent user message"},
+        {"role": "user", "content": "Recent user message 1"},
+        {"role": "user", "content": "Recent user message 2"},
+        {"role": "user", "content": "Recent user message 3"},
     ]
     
-    # Keep only the most recent message
-    result = trim_old_messages(messages, keep_recent=1)
+    result = trim_old_messages(messages)
     
     # System message should remain unchanged even though it's old
     system_content = result[0].get("content")
@@ -265,14 +266,22 @@ file2.txt
     assert FILES_SECTION_HEADING in assistant_content
     assert "file1.csv" in assistant_content
     
-    # Recent user message should remain unchanged
-    recent_content = result[3].get("content")
-    assert isinstance(recent_content, str)
-    assert recent_content == "Recent user message"
+    # Recent user messages should remain unchanged
+    recent_content_1 = result[3].get("content")
+    assert isinstance(recent_content_1, str)
+    assert recent_content_1 == "Recent user message 1"
+    
+    recent_content_2 = result[4].get("content")
+    assert isinstance(recent_content_2, str)
+    assert recent_content_2 == "Recent user message 2"
+    
+    recent_content_3 = result[5].get("content")
+    assert isinstance(recent_content_3, str)
+    assert recent_content_3 == "Recent user message 3"
 
 
 def test_trim_old_messages_preserves_recent_messages() -> None:
-    """Test that trim_old_messages preserves the most recent messages based on keep_recent parameter."""
+    """Test that trim_old_messages preserves the most recent messages based on MESSAGE_HISTORY_TRIM_THRESHOLD."""
     # Create test messages
     old_message_1 = f"""Old message 1.
 {FILES_SECTION_HEADING}
@@ -289,6 +298,10 @@ file3.csv
     recent_message_2 = f"""Recent message 2.
 {FILES_SECTION_HEADING}
 file4.csv
+"""
+    recent_message_3 = f"""Recent message 3.
+{FILES_SECTION_HEADING}
+file5.csv
 """
     
     # Create test messages with proper typing
@@ -297,10 +310,11 @@ file4.csv
         {"role": "user", "content": old_message_2},
         {"role": "user", "content": recent_message_1},
         {"role": "user", "content": recent_message_2},
+        {"role": "user", "content": recent_message_3},
     ]
     
-    # Keep the 2 most recent messages
-    result = trim_old_messages(messages, keep_recent=2)
+    # Test with MESSAGE_HISTORY_TRIM_THRESHOLD (3) - only the first 2 messages should be trimmed
+    result = trim_old_messages(messages)
     
     # Old messages should be trimmed
     old_content_1 = result[0].get("content")
@@ -325,24 +339,30 @@ file4.csv
     assert recent_content_2 == recent_message_2
     assert FILES_SECTION_HEADING in recent_content_2
     assert "file4.csv" in recent_content_2
+    
+    recent_content_3 = result[4].get("content")
+    assert isinstance(recent_content_3, str)
+    assert recent_content_3 == recent_message_3
+    assert FILES_SECTION_HEADING in recent_content_3
+    assert "file5.csv" in recent_content_3
 
 def test_trim_old_messages_empty_list() -> None:
     """Test that trim_old_messages handles empty message lists correctly."""
     messages: List[ChatCompletionMessageParam] = []
-    result = trim_old_messages(messages, keep_recent=2)
+    result = trim_old_messages(messages)
     assert result == []
 
 
-def test_trim_old_messages_fewer_than_keep_recent() -> None:
-    """Test that trim_old_messages doesn't modify messages if there are fewer than keep_recent."""
+def test_trim_old_messages_fewer_than_threshold() -> None:
+    """Test that trim_old_messages doesn't modify messages if there are fewer than MESSAGE_HISTORY_TRIM_THRESHOLD."""
     messages: List[ChatCompletionMessageParam] = [
         {"role": "user", "content": "User message 1"},
         {"role": "assistant", "content": "Assistant message 1"},
     ]
     
-    result = trim_old_messages(messages, keep_recent=3)
+    result = trim_old_messages(messages)
     
-    # Messages should remain unchanged
+    # Messages should remain unchanged since we have fewer than MESSAGE_HISTORY_TRIM_THRESHOLD (3) messages
     user_content = result[0].get("content")
     assert isinstance(user_content, str)
     assert user_content == "User message 1"
@@ -373,15 +393,17 @@ def test_trim_mixed_content_messages() -> None:
     })
     
     # Create sample message list with one old message (the mixed content)
-    # and one recent message (to not be trimmed)
+    # and enough recent messages to exceed MESSAGE_HISTORY_TRIM_THRESHOLD (3)
     message_list: List[ChatCompletionMessageParam] = [
         mixed_content_message,  # This should get trimmed
         {"role": "assistant", "content": "That's a chart showing data trends"},
-        {"role": "user", "content": "Can you explain more?"}  # Recent message, should not be trimmed
+        {"role": "user", "content": "Can you explain more?"},  # Recent message, should not be trimmed
+        {"role": "user", "content": "Another recent message"},  # Recent message, should not be trimmed
+        {"role": "user", "content": "Yet another recent message"}  # Recent message, should not be trimmed
     ]
     
     # Apply the trimming function
-    trimmed_messages = trim_old_messages(message_list, keep_recent=2)
+    trimmed_messages = trim_old_messages(message_list)
     
     # Verify that the first message has been trimmed properly
     assert trimmed_messages[0]["role"] == "user"
@@ -390,6 +412,8 @@ def test_trim_mixed_content_messages() -> None:
     # Verify that the recent messages are untouched
     assert trimmed_messages[1] == message_list[1]
     assert trimmed_messages[2] == message_list[2]
+    assert trimmed_messages[3] == message_list[3]
+    assert trimmed_messages[4] == message_list[4]
 
 
 def test_get_display_history_calls_update_last_interaction() -> None: