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

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -1,135 +1,138 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-import logging
 import os
 from anthropic.types import MessageParam
 from typing import List, Optional, Tuple, cast
-
-from mito_ai.logger import get_logger
-from mito_ai.streamlit_conversion.agent_utils import apply_patch_to_text, extract_todo_placeholders, fix_diff_headers
+from mito_ai.streamlit_conversion.agent_utils import apply_patch_to_text, extract_todo_placeholders, fix_diff_headers, 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.streamlit_system_prompt import streamlit_system_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, extract_unified_diff_blocks, parse_jupyter_notebook_to_extract_required_content
-from mito_ai.utils.anthropic_utils import stream_anthropic_completion_from_mito_server
+from mito_ai.streamlit_conversion.streamlit_utils import extract_code_blocks, create_app_file, extract_unified_diff_blocks, get_app_code_from_file, parse_jupyter_notebook_to_extract_required_content
 from mito_ai.completions.models import MessageType
 from mito_ai.utils.telemetry_utils import log_streamlit_app_creation_error, log_streamlit_app_creation_retry, log_streamlit_app_creation_success
 from mito_ai.streamlit_conversion.streamlit_utils import clean_directory_check
 
-STREAMLIT_AI_MODEL = "claude-3-5-haiku-latest"
-
-class StreamlitCodeGeneration:
-    @property
-    def log(self) -> logging.Logger:
-        """Use Mito AI logger."""
-        return get_logger()
-
-    async def get_response_from_agent(self, message_to_agent: List[MessageParam]) -> str:
-        """Gets the streaming response from the agent using the mito server"""
-        model = STREAMLIT_AI_MODEL
-        max_tokens = 8192 # 64_000
-        temperature = 0.2
-
-        self.log.info("Getting response from agent...")
-        accumulated_response = ""
-        async for stream_chunk in stream_anthropic_completion_from_mito_server(
-            model = model,
-            max_tokens = max_tokens,
-            temperature = temperature,
-            system = streamlit_system_prompt,
-            messages = message_to_agent,
-            stream=True,
-            message_type=MessageType.STREAMLIT_CONVERSION,
-            reply_fn=None,
-            message_id=""
-        ):
-            accumulated_response += stream_chunk
-        return accumulated_response
-
-    async def generate_streamlit_code(self, notebook: dict) -> str:
-        """Send a query to the agent, get its response and parse the code"""
-        
-        messages: List[MessageParam] = [
-            cast(MessageParam, {
-                "role": "user",
-                "content": [{
-                    "type": "text",
-                    "text": get_streamlit_app_creation_prompt(notebook)
-                }]
-            })
-        ]
-        
-        agent_response = await self.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 self.get_response_from_agent(todo_messages)
-            
-            # Apply the diff to the streamlit app
-            exctracted_diff = extract_unified_diff_blocks(todo_response)
-            fixed_diff = fix_diff_headers(exctracted_diff)
-            converted_code = apply_patch_to_text(converted_code, fixed_diff)
-                    
-        return converted_code
-
+def get_app_directory(notebook_path: str) -> str:
+    # Make sure the path is absolute if it is not already
+    absolute_notebook_path = os.path.abspath(notebook_path)
+    
+    # Get the directory of the notebook
+    app_directory = os.path.dirname(absolute_notebook_path)
+    return app_directory
 
-    async def correct_error_in_generation(self, 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] = [
+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": get_streamlit_error_correction_prompt(error, streamlit_app_code)
+                    "text": todo_prompt
                 }]
             })
         ]
-        agent_response = await self.get_response_from_agent(messages)
+        todo_response = await get_response_from_agent(todo_messages)
         
         # Apply the diff to the streamlit app
-        exctracted_diff = extract_unified_diff_blocks(agent_response)
-        
-        print(f"\n\nExtracted diff: {exctracted_diff}")
+        exctracted_diff = extract_unified_diff_blocks(todo_response)
         fixed_diff = fix_diff_headers(exctracted_diff)
-        streamlit_app_code = apply_patch_to_text(streamlit_app_code, fixed_diff)
-        
-        print("\n\nUpdated app code: ", streamlit_app_code)
+        converted_code = apply_patch_to_text(converted_code, fixed_diff)
+                
+    return converted_code
 
-        return streamlit_app_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)
+    exctracted_diff = extract_unified_diff_blocks(agent_response)
+    fixed_diff = fix_diff_headers(exctracted_diff)
+    print(fixed_diff)
+    converted_code = apply_patch_to_text(streamlit_app_code, fixed_diff)
+    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 diff to the streamlit app
+    exctracted_diff = extract_unified_diff_blocks(agent_response)
+    fixed_diff = fix_diff_headers(exctracted_diff)
+    streamlit_app_code = apply_patch_to_text(streamlit_app_code, fixed_diff)
 
-async def streamlit_handler(notebook_path: str) -> Tuple[bool, Optional[str], str]:
+    return streamlit_app_code
+
+async def streamlit_handler(notebook_path: str, edit_prompt: str = "") -> Tuple[bool, Optional[str], str]:
     """Handler function for streamlit code generation and validation"""
 
     clean_directory_check(notebook_path)
 
     notebook_code = parse_jupyter_notebook_to_extract_required_content(notebook_path)
-    streamlit_code_generator = StreamlitCodeGeneration()
-    streamlit_code = await streamlit_code_generator.generate_streamlit_code(notebook_code)
+    app_directory = get_app_directory(notebook_path)
     
+    if edit_prompt != "":
+        # If the user is editing an existing streamlit app, use the update function
+        streamlit_code = get_app_code_from_file(app_directory)
+        
+        if streamlit_code is None:
+            return False, '', "Error updating existing streamlit app because app.py file was not found."
+        
+        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 
     has_validation_error, errors = validate_app(streamlit_code, notebook_path)
     tries = 0
     while has_validation_error and tries < 5:
         for error in errors:
-            streamlit_code = await streamlit_code_generator.correct_error_in_generation(error, streamlit_code)
+            streamlit_code = await correct_error_in_generation(error, streamlit_code)
         
         has_validation_error, errors = validate_app(streamlit_code, notebook_path)
         
@@ -140,19 +143,13 @@ async def streamlit_handler(notebook_path: str) -> Tuple[bool, Optional[str], st
         tries+=1
 
     if has_validation_error:
-        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, error)
+        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, error, edit_prompt)
         return False, '', "Error generating streamlit code by agent"
     
-    # Convert to absolute path for directory calculation
-    absolute_notebook_path = notebook_path
-    if not (notebook_path.startswith('/') or (len(notebook_path) > 1 and notebook_path[1] == ':')):
-        absolute_notebook_path = os.path.join(os.getcwd(), notebook_path)
-    
-    app_directory = os.path.dirname(absolute_notebook_path)
+    # Finally, update the app.py file with the new code
     success_flag, app_path, message = create_app_file(app_directory, streamlit_code)
-    
     if not success_flag:
-        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, message)
+        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, message, edit_prompt)
     
-    log_streamlit_app_creation_success('mito_server_key', MessageType.STREAMLIT_CONVERSION)
+    log_streamlit_app_creation_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, edit_prompt)
     return success_flag, app_path, message

mito_ai/streamlit_conversion/streamlit_system_prompt.py

@@ -22,6 +22,7 @@ STREAMLIT IMPLEMENTATION GUIDELINES:
 - Include all text explanations and insights from markdown cells
 - Add interactive elements where beneficial (filters, selectors, etc.)
 - Ensure professional styling and layout suitable for executives
+- Just create the streamlit app code, do not include a _main_ function block. The file will be run directly using `streamlit run app.py`.
 
 CRITICAL REQUIREMENTS:
 1. **PRESERVE ALL CODE EXACTLY**: Every line of code, every data structure, every import must be included in full

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -4,7 +4,7 @@
 import re
 import json
 import os
-from typing import Dict, Optional, Tuple, Any
+from typing import Dict, List, Optional, Tuple, Any
 from pathlib import Path
 
 def extract_code_blocks(message_content: str) -> str:
@@ -25,7 +25,8 @@ def extract_code_blocks(message_content: str) -> str:
     matches = re.findall(pattern, message_content, re.DOTALL)
 
     # Concatenate with single newlines
-    return '\n'.join(matches)
+    result = '\n'.join(matches)
+    return result
 
 def extract_unified_diff_blocks(message_content: str) -> str:
     """
@@ -53,14 +54,23 @@ def create_app_file(app_directory: str, code: str) -> Tuple[bool, str, str]:
     """
     try:
         app_path = os.path.join(app_directory, "app.py")
-        with open(app_path, 'w') as f:
+
+        with open(app_path, 'w', encoding='utf-8') as f:
             f.write(code)
+        
         return True, app_path, f"Successfully created {app_directory}"
     except IOError as e:
         return False, '', f"Error creating file: {str(e)}"
     except Exception as e:
         return False, '', f"Unexpected error: {str(e)}"
     
+def get_app_code_from_file(app_directory: str) -> Optional[str]:
+    app_path = get_app_path(app_directory)
+    if app_path is None:
+        return None
+    with open(app_path, 'r', encoding='utf-8') as f:
+        return f.read()
+    
 
 def get_app_path(app_directory: str) -> Optional[str]:
     """
@@ -70,9 +80,8 @@ def get_app_path(app_directory: str) -> Optional[str]:
     if not os.path.exists(app_path):
         return None
     return app_path
-    
 
-def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Dict[str, Any]:
+def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> List[Dict[str, Any]]:
     """
     Read a Jupyter notebook and filter cells to keep only cell_type and source fields.
 
@@ -102,18 +111,15 @@ def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Di
             raise KeyError("Notebook does not contain 'cells' key")
 
         # Filter each cell to keep only cell_type and source
-        filtered_cells = []
+        filtered_cells: List[Dict[str, Any]] = []
         for cell in notebook_data['cells']:
-            filtered_cell = {
+            filtered_cell: Dict[str, Any] = {
                 'cell_type': cell.get('cell_type', ''),
                 'source': cell.get('source', [])
             }
             filtered_cells.append(filtered_cell)
-
-        # Update the notebook data with filtered cells
-        notebook_data['cells'] = filtered_cells
-
-        return notebook_data
+            
+        return filtered_cells
 
     except FileNotFoundError:
         raise FileNotFoundError(f"Notebook file not found: {notebook_path}")
@@ -139,8 +145,3 @@ def clean_directory_check(notebook_path: str) -> None:
 
     if not dir_path.exists():
         raise ValueError(f"Directory does not exist: {dir_path}")
-
-    file_count = len([f for f in dir_path.iterdir() if f.is_file()])
-    if file_count > 10:
-        raise ValueError(
-            f"Too many files in directory: 10 allowed but {file_count} present. Create a new directory and retry")

mito_ai/streamlit_conversion/validate_streamlit_app.py

@@ -23,42 +23,52 @@ 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:
-        self.temp_dir: Optional[str] = None
+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"""
+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
+        
+        original_cwd = os.getcwd()
         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
+            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. 
 
-    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)
-        
-        with change_working_directory(directory):
-            app_test = AppTest.from_string(app_code, default_timeout=30)
+        # 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
@@ -72,44 +82,38 @@ class StreamlitValidator:
                 return errors
             
             return None
+        finally:
+            # Clean up the temporary file
+            try:
+                os.unlink(temp_path)
+            except OSError:
+                pass  # File might already be deleted
 
-    def cleanup(self) -> None:
-        """Clean up the temporary files"""
-        if self.temp_dir and os.path.exists(self.temp_dir):
-            shutil.rmtree(self.temp_dir)
-            self.temp_dir = None
-
-    def _validate_app(self, 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})
+def check_for_errors(app_code: str, app_path: str) -> List[Dict[str, Any]]:
+    """Complete validation pipeline"""
+    errors: List[Dict[str, Any]] = []
 
-            runtime_errors = self.get_runtime_errors(app_code, app_path)
-            
-            print('Found Runtime Errors', runtime_errors)
-            
-            if runtime_errors:
-                errors.extend(runtime_errors)
-            
-        except Exception as e:
-            errors.append({'type': 'validation', 'details': str(e)})
+    try:
+        # Step 1: Check syntax
+        syntax_error = get_syntax_error(app_code)
+        if syntax_error:
+            errors.append({'type': 'syntax', 'details': syntax_error})
 
-        finally:
-            self.cleanup()
+        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,15 +82,17 @@ 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)
                 self.finish({"error": error_msg})

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}"