mito-ai 0.1.50__py3-none-any.whl

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -0,0 +1,144 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from anthropic.types import MessageParam
+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
+from mito_ai.streamlit_conversion.prompts.streamlit_finish_todo_prompt import get_finish_todo_prompt
+from mito_ai.streamlit_conversion.prompts.update_existing_app_prompt import get_update_existing_app_prompt
+from mito_ai.streamlit_conversion.validate_streamlit_app import validate_app
+from mito_ai.streamlit_conversion.streamlit_utils import extract_code_blocks, create_app_file, get_app_code_from_file, parse_jupyter_notebook_to_extract_required_content
+from mito_ai.streamlit_conversion.search_replace_utils import extract_search_replace_blocks, apply_search_replace
+from mito_ai.completions.models import MessageType
+from mito_ai.utils.error_classes import StreamlitConversionError
+from mito_ai.utils.telemetry_utils import log_streamlit_app_validation_retry, log_streamlit_app_conversion_success
+from mito_ai.path_utils import AbsoluteNotebookPath, AppFileName, get_absolute_notebook_dir_path, get_absolute_app_path, get_app_file_name
+
+async def generate_new_streamlit_code(notebook: List[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": prompt_text
+            }]
+        })
+    ]
+    agent_response = await get_response_from_agent(messages)
+    converted_code = extract_code_blocks(agent_response)
+    
+    # Extract the TODOs from the agent's response
+    todo_placeholders = extract_todo_placeholders(agent_response)
+    
+    for todo_placeholder in todo_placeholders:
+        print(f"Processing AI TODO: {todo_placeholder}")
+        todo_prompt = get_finish_todo_prompt(notebook, converted_code, todo_placeholder)
+        todo_messages: List[MessageParam] = [
+            cast(MessageParam, {
+                "role": "user",
+                "content": [{
+                    "type": "text",
+                    "text": todo_prompt
+                }]
+            })
+        ]
+        todo_response = await get_response_from_agent(todo_messages)
+        
+        # Apply the search/replace to the streamlit app
+        search_replace_pairs = extract_search_replace_blocks(todo_response)
+        converted_code = apply_search_replace(converted_code, search_replace_pairs)
+                
+    return converted_code
+
+
+async def update_existing_streamlit_code(notebook: List[dict], streamlit_app_code: str, edit_prompt: str) -> str:
+    """Send a query to the agent, get its response and parse the code"""
+    prompt_text = get_update_existing_app_prompt(notebook, streamlit_app_code, edit_prompt)
+    
+    messages: List[MessageParam] = [
+        cast(MessageParam, {
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": prompt_text
+            }]
+        })
+    ]
+    
+    agent_response = await get_response_from_agent(messages)
+    print(f"[Mito AI Search/Replace Tool]:\n {agent_response}")
+
+    # Apply the search/replace to the streamlit app
+    search_replace_pairs = extract_search_replace_blocks(agent_response)
+    converted_code = apply_search_replace(streamlit_app_code, search_replace_pairs)
+    print(f"[Mito AI Search/Replace Tool]\nConverted code\n: {converted_code}")
+    return converted_code
+
+
+async def correct_error_in_generation(error: str, streamlit_app_code: str) -> str:
+    """If errors are present, send it back to the agent to get corrections in code"""
+    messages: List[MessageParam] = [
+        cast(MessageParam, {
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": get_streamlit_error_correction_prompt(error, streamlit_app_code)
+            }]
+        })
+    ]
+    agent_response = await get_response_from_agent(messages)
+    
+    # Apply the search/replace to the streamlit app
+    search_replace_pairs = extract_search_replace_blocks(agent_response)
+    streamlit_app_code = apply_search_replace(streamlit_app_code, search_replace_pairs)
+
+    return streamlit_app_code
+
+async def streamlit_handler(notebook_path: AbsoluteNotebookPath, app_file_name: AppFileName, edit_prompt: str = "") -> None:
+    """Handler function for streamlit code generation and validation"""
+
+    # Convert to absolute path for consistent handling
+    notebook_code = parse_jupyter_notebook_to_extract_required_content(notebook_path)
+    app_directory = get_absolute_notebook_dir_path(notebook_path)
+    app_path = get_absolute_app_path(app_directory, app_file_name)
+    
+    if edit_prompt != "":
+        # If the user is editing an existing streamlit app, use the update function
+        streamlit_code = get_app_code_from_file(app_path)
+        
+        if streamlit_code is None:
+            raise StreamlitConversionError("Error updating existing streamlit app because app.py file was not found.", 404)
+        
+        streamlit_code = await update_existing_streamlit_code(notebook_code, streamlit_code, edit_prompt)
+    else:
+        # Otherwise generate a new streamlit app
+        streamlit_code = await generate_new_streamlit_code(notebook_code)
+       
+    # Then, after creating/updating the app, validate that the new code runs 
+    errors = validate_app(streamlit_code, notebook_path)
+    tries = 0
+    while len(errors)>0 and tries < 5:
+        for error in errors:
+            streamlit_code = await correct_error_in_generation(error, streamlit_code)
+        
+        errors = validate_app(streamlit_code, notebook_path)
+        
+        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 len(errors)>0:
+        final_errors = ', '.join(errors)
+        raise StreamlitConversionError(f"Streamlit agent failed generating code after max retries. Errors: {final_errors}", 500)
+    
+    # Finally, update the app.py file with the new code
+    create_app_file(app_path, streamlit_code)
+    log_streamlit_app_conversion_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, edit_prompt)

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -0,0 +1,85 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import re
+import json
+from typing import Dict, List, Optional, Any
+from mito_ai.path_utils import AbsoluteAppPath, AbsoluteNotebookPath
+from mito_ai.utils.error_classes import StreamlitConversionError
+
+
+def extract_code_blocks(message_content: str) -> str:
+    """
+    Extract all code blocks from Claude's response.
+
+    Args:
+        message_content (str): The actual content from the agent's response
+
+    Returns:
+        str: Removes the ```python``` part to be able to parse the code
+    """
+    if "```python" not in message_content:
+        return message_content
+
+    # Use regex to find all Python code blocks
+    pattern = r'```python\n(.*?)```'
+    matches = re.findall(pattern, message_content, re.DOTALL)
+
+    # Concatenate with single newlines
+    result = '\n'.join(matches)
+    return result
+
+def create_app_file(app_path: AbsoluteAppPath, code: str) -> None:
+    """
+    Create .py file and write code to it with error handling
+    """
+    try:
+        with open(app_path, 'w', encoding='utf-8') as f:
+            f.write(code)
+    except IOError as e:
+        raise StreamlitConversionError(f"Error creating app file: {str(e)}", 500)
+    
+def get_app_code_from_file(app_path: AbsoluteAppPath) -> Optional[str]:
+    with open(app_path, 'r', encoding='utf-8') as f:
+        return f.read()
+
+def parse_jupyter_notebook_to_extract_required_content(notebook_path: AbsoluteNotebookPath) -> List[Dict[str, Any]]:
+    """
+    Read a Jupyter notebook and filter cells to keep only cell_type and source fields.
+
+    Args:
+        notebook_path: Absolute path to the .ipynb file
+
+    Returns:
+        dict: Filtered notebook dictionary with only cell_type and source in cells
+
+    Raises:
+        FileNotFoundError: If the notebook file doesn't exist
+        json.JSONDecodeError: If the file is not valid JSON
+        KeyError: If the notebook doesn't have the expected structure
+    """
+    
+    try:
+        # Read the notebook file
+        with open(notebook_path, 'r', encoding='utf-8') as f:
+            notebook_data: Dict[str, Any] = json.load(f)
+
+        # Check if 'cells' key exists
+        if 'cells' not in notebook_data:
+            raise StreamlitConversionError("Notebook does not contain 'cells' key", 400)
+
+        # Filter each cell to keep only cell_type and source
+        filtered_cells: List[Dict[str, Any]] = []
+        for cell in notebook_data['cells']:
+            filtered_cell: Dict[str, Any] = {
+                'cell_type': cell.get('cell_type', ''),
+                'source': cell.get('source', [])
+            }
+            filtered_cells.append(filtered_cell)
+            
+        return filtered_cells
+
+    except FileNotFoundError:
+        raise StreamlitConversionError(f"Notebook file not found: {notebook_path}", 404)
+    except json.JSONDecodeError as e:
+        raise StreamlitConversionError(f"Invalid JSON in notebook file: {str(e)}", 400)

mito_ai/streamlit_conversion/validate_streamlit_app.py

@@ -0,0 +1,105 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import tempfile
+import traceback
+import ast
+import warnings
+from typing import List, Tuple, Optional, Dict, Any, Generator
+from streamlit.testing.v1 import AppTest
+from contextlib import contextmanager
+from mito_ai.path_utils import AbsoluteNotebookPath, get_absolute_notebook_dir_path
+
+warnings.filterwarnings("ignore", message=".*bare mode.*")
+
+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_runtime_errors(app_code: str, app_path: AbsoluteNotebookPath) -> Optional[List[Dict[str, Any]]]:
+    """Start the Streamlit app in a subprocess"""  
+    
+    directory = get_absolute_notebook_dir_path(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)
+    
+    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
+
+        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 check_for_errors(app_code: str, app_path: AbsoluteNotebookPath) -> List[Dict[str, Any]]:
+    """Complete validation pipeline"""
+    errors: List[Dict[str, Any]] = []
+
+    try:
+        # Step 1: Check syntax
+        syntax_error = get_syntax_error(app_code)
+        if syntax_error:
+            errors.append({'type': 'syntax', 'details': syntax_error})
+
+        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
+
+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)
+    stringified_errors = [str(error) for error in errors]
+    return stringified_errors

mito_ai/streamlit_preview/__init__.py

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

mito_ai/streamlit_preview/handlers.py

@@ -0,0 +1,111 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import uuid
+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 StreamlitPreviewManager
+from mito_ai.path_utils import get_absolute_notebook_dir_path, get_absolute_notebook_path, get_absolute_app_path, does_app_path_exist, get_app_file_name
+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
+
+
+class StreamlitPreviewHandler(APIHandler):
+    """REST handler for streamlit preview operations."""
+
+    def initialize(self) -> None:
+        """Initialize the handler."""
+        self.preview_manager = StreamlitPreviewManager()
+
+    @tornado.web.authenticated
+    async def post(self) -> None:
+        """Start a new streamlit preview."""
+        try:
+            # Parse and validate request
+            body = self.get_json_body()
+            notebook_path, notebook_id, force_recreate, edit_prompt = validate_request_body(body)
+
+            # Ensure app exists
+            absolute_notebook_path = get_absolute_notebook_path(notebook_path)
+            absolute_notebook_dir_path = get_absolute_notebook_dir_path(absolute_notebook_path)
+            app_file_name = get_app_file_name(notebook_id)
+            absolute_app_path = get_absolute_app_path(absolute_notebook_dir_path, app_file_name)
+            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, app_file_name, 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())
+            port = self.preview_manager.start_streamlit_preview(absolute_notebook_dir_path, app_file_name, preview_id)
+
+            # Return success response
+            self.finish({
+                "type": 'success',
+                "id": preview_id, 
+                "port": port, 
+                "url": f"http://localhost:{port}"
+            })
+            log_streamlit_app_preview_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, edit_prompt)
+
+        except StreamlitConversionError as e:
+            print(e)
+            self.set_status(e.error_code)
+            error_message = str(e)
+            formatted_traceback = traceback.format_exc()
+            self.finish({"error": error_message})
+            log_streamlit_app_conversion_error(
+                'mito_server_key', 
+                MessageType.STREAMLIT_CONVERSION, 
+                error_message, 
+                formatted_traceback,
+                edit_prompt,
+            )
+        except StreamlitPreviewError as e:
+            print(e)
+            error_message = str(e)
+            formatted_traceback = traceback.format_exc()
+            self.set_status(e.error_code)
+            self.finish({"error": error_message})
+            log_streamlit_app_preview_failure('mito_server_key', MessageType.STREAMLIT_CONVERSION, error_message, formatted_traceback, edit_prompt)
+        except Exception as e:
+            print(f"Exception in streamlit preview handler: {e}")
+            self.set_status(500)
+            error_message = str(e)
+            formatted_traceback = traceback.format_exc()
+            self.finish({"error": error_message})
+            log_streamlit_app_preview_failure('mito_server_key', MessageType.STREAMLIT_CONVERSION, error_message, formatted_traceback, "")
+
+    @tornado.web.authenticated
+    def delete(self, preview_id: str) -> None:
+        """Stop a streamlit preview."""
+        try:
+            if not preview_id:
+                self.set_status(400)
+                self.finish({"error": "Missing preview_id parameter"})
+                return
+
+            # Stop the preview
+            stopped = self.preview_manager.stop_preview(preview_id)
+
+            if stopped:
+                self.set_status(204)  # No content
+            else:
+                self.set_status(404)
+                self.finish({"error": f"Preview {preview_id} not found"})
+
+        except Exception as e:
+            self.set_status(500)
+            self.finish({"error": f"Internal server error: {str(e)}"})

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/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+