mito-ai 0.1.33__py3-none-any.whl → 0.1.49__py3-none-any.whl

mito_ai/streamlit_preview/manager.py

@@ -0,0 +1,152 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import socket
+import subprocess
+import time
+import threading
+import requests
+from typing import Dict, Optional, Tuple
+from dataclasses import dataclass
+from mito_ai.logger import get_logger
+from mito_ai.path_utils import AbsoluteNotebookDirPath, AppFileName
+from mito_ai.utils.error_classes import StreamlitPreviewError
+
+
+@dataclass
+class PreviewProcess:
+    """Data class to track a streamlit preview process."""
+    proc: subprocess.Popen
+    port: int
+
+
+class StreamlitPreviewManager:
+    """Manages streamlit preview processes and their lifecycle."""
+    
+    def __init__(self) -> None:
+        self._previews: Dict[str, PreviewProcess] = {}
+        self._lock = threading.Lock()
+        self.log = get_logger()
+        
+    def get_free_port(self) -> int:
+        """Get a free port for streamlit to use."""
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(('', 0))
+            s.listen(1)
+            port = int(s.getsockname()[1])
+        
+        return port
+    
+    def start_streamlit_preview(self, app_directory: AbsoluteNotebookDirPath, app_file_name: AppFileName, preview_id: str) -> int:
+        """Start a streamlit preview process.
+        
+        Args:
+            app_code: The streamlit app code to run
+            preview_id: Unique identifier for this preview
+            
+        Returns:
+            Tuple of (success, message, port)
+        """
+        
+        try:
+            
+            # Get free port
+            port = self.get_free_port()
+            
+            # Start streamlit process
+            cmd = [
+                "streamlit", "run", app_file_name,
+                "--server.port", str(port),
+                "--server.headless", "true",
+                "--server.address", "localhost",
+                "--server.enableXsrfProtection", "false",
+                "--server.runOnSave", "true",  # auto-reload when app is saved
+                "--logger.level", "error"
+            ]
+            
+            # TODO: Security considerations for production:
+            # - Consider enabling XSRF protection if needed, but we might already get this with the APIHandler?
+            # - Add authentication headers to streamlit
+            
+            proc = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                cwd=app_directory
+            )
+            
+            # Wait for app to be ready
+            ready = self._wait_for_app_ready(port)
+            if not ready:
+                proc.terminate()
+                proc.wait()
+                raise StreamlitPreviewError("Streamlit app failed to start as app is not ready", 500)
+            
+            # Register the process
+            with self._lock:
+                self._previews[preview_id] = PreviewProcess(
+                    proc=proc,
+                    port=port,
+                )
+            
+            self.log.info(f"Started streamlit preview {preview_id} on port {port}")
+            return port
+            
+        except Exception as e:
+            self.log.error(f"Error starting streamlit preview: {e}")
+            raise StreamlitPreviewError(f"Failed to start preview: {str(e)}", 500)
+    
+    def _wait_for_app_ready(self, port: int, timeout: int = 30) -> bool:
+        """Wait for streamlit app to be ready on the given port."""
+        start_time = time.time()
+        
+        while time.time() - start_time < timeout:
+            try:
+                response = requests.get(f"http://localhost:{port}", timeout=5)
+                if response.status_code == 200:
+                    return True
+            except requests.exceptions.RequestException as e:
+                self.log.info(f"Waiting for app to be ready...")
+                pass
+            
+            time.sleep(1)
+        
+        return False
+    
+    def stop_preview(self, preview_id: str) -> bool:
+        """Stop a streamlit preview process.
+        
+        Args:
+            preview_id: The preview ID to stop
+            
+        Returns:
+            True if stopped successfully, False if not found
+        """
+        self.log.info(f"Stopping preview {preview_id}")
+        with self._lock:
+            if preview_id not in self._previews:
+                return False
+            
+            preview = self._previews[preview_id]
+            
+            # Terminate process
+            try:
+                preview.proc.terminate()
+                preview.proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                preview.proc.kill()
+                preview.proc.wait()
+            except Exception as e:
+                self.log.error(f"Error terminating process {preview_id}: {e}")
+            
+            # Remove from registry
+            del self._previews[preview_id]
+            
+            self.log.info(f"Stopped streamlit preview {preview_id}")
+            return True
+    
+    def get_preview(self, preview_id: str) -> Optional[PreviewProcess]:
+        """Get a preview process by ID."""
+        with self._lock:
+            return self._previews.get(preview_id)

mito_ai/streamlit_preview/urls.py

@@ -0,0 +1,22 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Any, List, Tuple
+from jupyter_server.utils import url_path_join
+from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
+
+def get_streamlit_preview_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all streamlit preview related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    
+    return [
+        (url_path_join(BASE_URL, "streamlit-preview"), StreamlitPreviewHandler, {}),
+        (url_path_join(BASE_URL, "streamlit-preview/(.+)"), StreamlitPreviewHandler, {}),
+    ]

mito_ai/streamlit_preview/utils.py

@@ -0,0 +1,29 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Tuple, Optional
+from mito_ai.utils.error_classes import StreamlitPreviewError
+
+
+def validate_request_body(body: Optional[dict]) -> Tuple[str, str, bool, str]:
+    """Validate the request body and extract notebook_path and force_recreate."""
+    if body is None:
+        raise StreamlitPreviewError("Invalid or missing JSON body", 400)
+
+    notebook_path = body.get("notebook_path")
+    if not notebook_path:
+        raise StreamlitPreviewError("Missing notebook_path parameter", 400)
+    
+    notebook_id = body.get("notebook_id")
+    if not notebook_id:
+        raise StreamlitPreviewError("Missing notebook_id parameter", 400)
+
+    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, notebook_id, force_recreate, edit_prompt

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/completions/completion_handlers_utils_test.py

@@ -0,0 +1,190 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import base64
+import os
+import tempfile
+from contextlib import contextmanager
+from mito_ai.completions.completion_handlers.utils import (
+    create_ai_optimized_message,
+    extract_and_encode_images_from_additional_context,
+)
+
+
+@contextmanager
+def temporary_image_file(suffix=".png", content=b"fake_image_data"):
+    """Context manager that creates a temporary image file for testing."""
+    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as temp_file:
+        temp_file.write(content)
+        temp_file_path = temp_file.name
+
+    try:
+        yield temp_file_path
+    finally:
+        # Clean up the temporary file
+        os.unlink(temp_file_path)
+
+
+def test_text_only_message():
+    """Test scenario where the user only inputs text"""
+    result = create_ai_optimized_message("Hello world")
+
+    assert result["role"] == "user"
+    assert result["content"] == "Hello world"
+
+
+def test_message_with_uploaded_image():
+    """Test scenario where the user uploads an image"""
+    with temporary_image_file() as temp_file_path:
+        result = create_ai_optimized_message(
+            text="Analyze this",
+            additional_context=[{"type": "image/png", "value": temp_file_path}],
+        )
+
+        assert result["role"] == "user"
+        assert isinstance(result["content"], list)
+        assert result["content"][0]["type"] == "text"
+        assert result["content"][1]["type"] == "image_url"
+
+
+def test_message_with_multiple_uploaded_images():
+    """Test scenario where the user uploads multiple images"""
+    with temporary_image_file(suffix=".png", content=b"image1_data") as temp_file1:
+        with temporary_image_file(suffix=".jpg", content=b"image2_data") as temp_file2:
+            result = create_ai_optimized_message(
+                text="Analyze these images",
+                additional_context=[
+                    {"type": "image/png", "value": temp_file1},
+                    {"type": "image/jpeg", "value": temp_file2},
+                ],
+            )
+
+            assert result["role"] == "user"
+            assert isinstance(result["content"], list)
+            assert len(result["content"]) == 3  # text + 2 images
+            assert result["content"][0]["type"] == "text"
+            assert result["content"][0]["text"] == "Analyze these images"
+            assert result["content"][1]["type"] == "image_url"
+            assert result["content"][2]["type"] == "image_url"
+            
+            # Verify the image URLs are properly formatted
+            assert result["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
+            assert result["content"][2]["image_url"]["url"].startswith("data:image/jpeg;base64,")
+
+
+def test_message_with_active_cell_output():
+    """Test scenario where the active cell has an output"""
+    result = create_ai_optimized_message(
+        text="Analyze this", base64EncodedActiveCellOutput="cell_output_data"
+    )
+
+    assert result["role"] == "user"
+    assert isinstance(result["content"], list)
+    assert result["content"][0]["type"] == "text"
+    assert result["content"][1]["type"] == "image_url"
+
+
+def test_message_with_uploaded_image_and_active_cell_output():
+    """Test scenario where the user uploads an image and the active cell has an output"""
+    with temporary_image_file() as temp_file_path:
+        result = create_ai_optimized_message(
+            text="Analyze this",
+            additional_context=[{"type": "image/png", "value": temp_file_path}],
+            base64EncodedActiveCellOutput="cell_output_data",
+        )
+
+        assert result["role"] == "user"
+        assert isinstance(result["content"], list)
+        assert result["content"][0]["type"] == "text"
+        assert result["content"][1]["type"] == "image_url"
+        assert result["content"][2]["type"] == "image_url"
+
+
+def test_extract_and_encode_images_from_additional_context_valid_image():
+    """Test extracting and encoding a valid image file"""
+    with temporary_image_file() as temp_file_path:
+        additional_context = [{"type": "image/png", "value": temp_file_path}]
+
+        encoded_images = extract_and_encode_images_from_additional_context(
+            additional_context
+        )
+
+        assert len(encoded_images) == 1
+        assert encoded_images[0].startswith("data:image/png;base64,")
+        # Verify it's valid base64 by checking it can be decoded
+        base64_data = encoded_images[0].split(",")[1]
+        decoded_data = base64.b64decode(base64_data)
+        assert decoded_data == b"fake_image_data"
+
+
+def test_extract_and_encode_images_from_additional_context_multiple_images():
+    """Test extracting and encoding multiple image files"""
+    with temporary_image_file(suffix=".png", content=b"image1_data") as temp_file1:
+        with temporary_image_file(suffix=".jpg", content=b"image2_data") as temp_file2:
+            additional_context = [
+                {"type": "image/png", "value": temp_file1},
+                {"type": "image/jpeg", "value": temp_file2},
+            ]
+
+            encoded_images = extract_and_encode_images_from_additional_context(
+                additional_context
+            )
+
+            assert len(encoded_images) == 2
+            assert encoded_images[0].startswith("data:image/png;base64,")
+            assert encoded_images[1].startswith("data:image/jpeg;base64,")
+
+
+def test_extract_and_encode_images_from_additional_context_invalid_file():
+    """Test handling of invalid/non-existent image files"""
+    additional_context = [{"type": "image/png", "value": "non_existent_file.png"}]
+
+    encoded_images = extract_and_encode_images_from_additional_context(
+        additional_context
+    )
+
+    assert len(encoded_images) == 0
+
+
+def test_extract_and_encode_images_from_additional_context_non_image_types():
+    """Test that non-image types are ignored"""
+    with temporary_image_file(suffix=".txt", content=b"text_data") as temp_file:
+        additional_context = [
+            {"type": "text/plain", "value": temp_file},
+            {"type": "application/pdf", "value": "document.pdf"},
+        ]
+
+        encoded_images = extract_and_encode_images_from_additional_context(
+            additional_context
+        )
+
+        assert len(encoded_images) == 0
+
+
+def test_extract_and_encode_images_from_additional_context_mixed_types():
+    """Test handling of mixed image and non-image types"""
+    with temporary_image_file() as temp_image_file:
+        additional_context = [
+            {"type": "image/png", "value": temp_image_file},
+            {"type": "text/plain", "value": "document.txt"},
+            {"type": "image/jpeg", "value": "non_existent.jpg"},
+        ]
+
+        encoded_images = extract_and_encode_images_from_additional_context(
+            additional_context
+        )
+
+        # Should only have the valid PNG image
+        assert len(encoded_images) == 1
+        assert encoded_images[0].startswith("data:image/png;base64,")
+
+
+def test_extract_and_encode_images_from_additional_context_empty():
+    """Test handling of empty or None additional_context"""
+    # Test with None
+    encoded_images = extract_and_encode_images_from_additional_context(None)
+    assert len(encoded_images) == 0
+
+    # Test with empty list
+    encoded_images = extract_and_encode_images_from_additional_context([])
+    assert len(encoded_images) == 0