mito-ai 0.1.38__py3-none-any.whl → 0.1.40__py3-none-any.whl

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -4,31 +4,24 @@
 import logging
 import os
 from anthropic.types import MessageParam
-from typing import List, Optional, Tuple, cast
+from typing import List, Optional, Tuple, cast, Union
 
 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.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.validate_and_run_streamlit_code import streamlit_code_validator
-from mito_ai.streamlit_conversion.streamlit_utils import extract_code_blocks, create_app_file, parse_jupyter_notebook_to_extract_required_content
+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.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:
-    def __init__(self, notebook: dict) -> None:
-
-        self.messages: List[MessageParam] = [
-            cast(MessageParam, {
-                "role": "user",
-                "content": [{
-                    "type": "text",
-                    "text": f"Here is my jupyter notebook content that I want to convert into a Streamlit dashboard - {notebook}"
-                }]
-            })
-        ]
-
     @property
     def log(self) -> logging.Logger:
         """Use Mito AI logger."""
@@ -56,57 +49,89 @@ class StreamlitCodeGeneration:
             accumulated_response += stream_chunk
         return accumulated_response
 
-    def add_agent_response_to_context(self, agent_response: str) -> None:
-        """Add the agent's response to the history"""
-        self.messages.append(
+    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": "assistant",
+                "role": "user",
                 "content": [{
                     "type": "text",
-                    "text": agent_response
+                    "text": get_streamlit_app_creation_prompt(notebook)
                 }]
             })
-        )
-
-    async def generate_streamlit_code(self) -> str:
-        """Send a query to the agent, get its response and parse the code"""
-        agent_response = await self.get_response_from_agent(self.messages)
+        ]
+        
+        agent_response = await self.get_response_from_agent(messages)
 
         converted_code = extract_code_blocks(agent_response)
-        self.add_agent_response_to_context(converted_code)
+        
+        # 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
 
 
-    async def correct_error_in_generation(self, error: str) -> str:
+    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"""
-        self.messages.append(
+        messages: List[MessageParam] = [
             cast(MessageParam, {
                 "role": "user",
                 "content": [{
                     "type": "text",
-                    "text": f"When I run the streamlit app code, I get the following error: {error}\nPlease return the FULL Streamlit app code with the error corrected"
+                    "text": get_streamlit_error_correction_prompt(error, streamlit_app_code)
                 }]
             })
-        )
-        agent_response = await self.get_response_from_agent(self.messages)
-        converted_code = extract_code_blocks(agent_response)
-        self.add_agent_response_to_context(converted_code)
+        ]
+        agent_response = await self.get_response_from_agent(messages)
+        
+        # Apply the diff to the streamlit app
+        exctracted_diff = extract_unified_diff_blocks(agent_response)
+        
+        print(f"\n\nExtracted diff: {exctracted_diff}")
+        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)
 
-        return converted_code
+        return streamlit_app_code
 
 
 async def streamlit_handler(notebook_path: 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(notebook_code)
-    streamlit_code = await streamlit_code_generator.generate_streamlit_code()
-    has_validation_error, error = streamlit_code_validator(streamlit_code)
-    
+    streamlit_code_generator = StreamlitCodeGeneration()
+    streamlit_code = await streamlit_code_generator.generate_streamlit_code(notebook_code)
     
+    has_validation_error, errors = validate_app(streamlit_code, notebook_path)
     tries = 0
     while has_validation_error and tries < 5:
-        streamlit_code = await streamlit_code_generator.correct_error_in_generation(error)
-        has_validation_error, error = streamlit_code_validator(streamlit_code)
+        for error in errors:
+            streamlit_code = await streamlit_code_generator.correct_error_in_generation(error, streamlit_code)
+        
+        has_validation_error, errors = validate_app(streamlit_code, notebook_path)
         
         if has_validation_error:
             # TODO: We can't easily get the key type here, so for the beta release
@@ -116,9 +141,14 @@ async def streamlit_handler(notebook_path: str) -> Tuple[bool, Optional[str], st
 
     if has_validation_error:
         log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, error)
-        return False, None, "Error generating streamlit code by agent"
+        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(notebook_path)
+    app_directory = os.path.dirname(absolute_notebook_path)
     success_flag, app_path, message = create_app_file(app_directory, streamlit_code)
     
     if not success_flag:

mito_ai/streamlit_conversion/streamlit_system_prompt.py

@@ -1,7 +1,7 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-streamlit_system_prompt = """You are a senior data scientist and Streamlit expert specializing in converting Jupyter notebooks into professional dashboard applications.
+streamlit_system_prompt = """You are a code conversion specialist who converts Jupyter notebooks into Streamlit applications with ABSOLUTE FIDELITY.
 
 ROLE AND EXPERTISE:
 - Expert in Python, Jupyter notebooks, Streamlit, and data visualization
@@ -17,17 +17,25 @@ TASK REQUIREMENTS:
 STREAMLIT IMPLEMENTATION GUIDELINES:
 - Use appropriate Streamlit components (st.title, st.header, st.subheader, st.markdown, etc.)
 - Display all visualizations using st.pyplot(), st.plotly_chart(), or st.altair_chart() as appropriate
+- Do not convert database connections into Streamlit's secret.toml format. If the user inlined their database credentials, are importing from an environment variable, or reading from a connections file, assume that same approach will work in the streamlit app.
 - Show dataframes and tables using st.dataframe() or st.table()
 - 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
 
-CODE STRUCTURE:
-- Generate a complete, runnable app.py file
-- Include all necessary imports
-- Handle data loading and processing
-- Organize content with clear sections and headers
-- Include comments explaining key sections
+CRITICAL REQUIREMENTS:
+1. **PRESERVE ALL CODE EXACTLY**: Every line of code, every data structure, every import must be included in full
+2. **NO PLACEHOLDERS**: Never use comments like "# Add more data here" or "# Fill in the rest"
+3. **NO SIMPLIFICATION**: Do not replace actual data with sample data or hardcoded examples
+4. **COMPLETE DATA STRUCTURES**: If a notebook has a 1000-line dictionary, include all 1000 lines
+5. **PRESERVE DATA LOADING**: If the notebook reads from files, the Streamlit app must read from the same files
+6. **NO IMPROVIZAITION**: Do not provide your own interpretations of the analysis. Just convert the existing analysis into a streamlit app.
+
+STYLE GUIDELINES: 
+- Create a professional, executive-friendly dashboard
+- If there are variables in the notebook that the streamlit app viewer would likely want to configure, then use the appropriate streamlit component to allow them to do so. For examples, if the notebook has a variable called "start_date" and "end_date", then use the st.date_input component to allow the user to select the start and end dates.
+- Do not use emojis unless they are in the notebook already
+- Do not modify the graphs or analysis. If the notebook has a graph, use the same graph in the streamlit app.
 - Always include the following code at the top of the file so the user does not use the wrong deploy button
 ```python
 st.markdown(\"\"\"
@@ -41,13 +49,7 @@ st.markdown(\"\"\"
 ```
 
 OUTPUT FORMAT:
-- Provide the complete app.py file code
-- Ensure all notebook outputs are faithfully reproduced
-- Make the dashboard professional and presentation-ready
-- Focus on clarity and executive-level communication
-- Don't give extra explanations, just give the python code 
-- Do NOT add emojis 
-- Do NOT modify the graphs or analysis
-- Do NOT provide your own interpretations for the analysis
-
-Remember: The goal is to transform technical analysis into a polished, interactive/visually appealing dashboard that executives can easily understand and navigate."""
+- Output the complete, runnable app.py file.
+- Do not output any extra text, just give the python code. 
+
+"""

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -5,6 +5,7 @@ import re
 import json
 import os
 from typing import Dict, Optional, Tuple, Any
+from pathlib import Path
 
 def extract_code_blocks(message_content: str) -> str:
     """
@@ -19,7 +20,6 @@ def extract_code_blocks(message_content: str) -> str:
     if "```python" not in message_content:
         return message_content
 
-    # return message_content.split('```python\n')[1].split('\n```')[0]
     # Use regex to find all Python code blocks
     pattern = r'```python\n(.*?)```'
     matches = re.findall(pattern, message_content, re.DOTALL)
@@ -27,8 +27,19 @@ def extract_code_blocks(message_content: str) -> str:
     # Concatenate with single newlines
     return '\n'.join(matches)
 
+def extract_unified_diff_blocks(message_content: str) -> str:
+    """
+    Extract all unified_diff blocks from Claude's response.
+    """
+    if "```unified_diff" not in message_content:
+        return message_content
+    
+    pattern = r'```unified_diff\n(.*?)```'
+    matches = re.findall(pattern, message_content, re.DOTALL)
+    return '\n'.join(matches)
+
 
-def create_app_file(app_directory: str, code: str) -> Tuple[bool, Optional[str], str]:
+def create_app_file(app_directory: str, code: str) -> Tuple[bool, str, str]:
     """
     Create app.py file and write code to it with error handling
 
@@ -46,9 +57,9 @@ def create_app_file(app_directory: str, code: str) -> Tuple[bool, Optional[str],
             f.write(code)
         return True, app_path, f"Successfully created {app_directory}"
     except IOError as e:
-        return False, None, f"Error creating file: {str(e)}"
+        return False, '', f"Error creating file: {str(e)}"
     except Exception as e:
-        return False, None, f"Unexpected error: {str(e)}"
+        return False, '', f"Unexpected error: {str(e)}"
 
 
 def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Dict[str, Any]:
@@ -56,7 +67,7 @@ def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Di
     Read a Jupyter notebook and filter cells to keep only cell_type and source fields.
 
     Args:
-        notebook_path (str): Absolute path to the .ipynb file
+        notebook_path (str): Path to the .ipynb file (can be relative or absolute)
 
     Returns:
         dict: Filtered notebook dictionary with only cell_type and source in cells
@@ -66,6 +77,11 @@ def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Di
         json.JSONDecodeError: If the file is not valid JSON
         KeyError: If the notebook doesn't have the expected structure
     """
+    # Convert to absolute path if it's not already absolute
+    # Handle both Unix-style absolute paths (starting with /) and Windows-style absolute paths
+    if not (notebook_path.startswith('/') or (len(notebook_path) > 1 and notebook_path[1] == ':')):
+        notebook_path = os.path.join(os.getcwd(), notebook_path)
+    
     try:
         # Read the notebook file
         with open(notebook_path, 'r', encoding='utf-8') as f:
@@ -96,3 +112,25 @@ def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Di
         raise json.JSONDecodeError(f"Invalid JSON in notebook file: {str(e)}", e.doc if hasattr(e, 'doc') else '', e.pos if hasattr(e, 'pos') else 0)
     except Exception as e:
         raise Exception(f"Error processing notebook: {str(e)}")
+
+
+def resolve_notebook_path(notebook_path:str) -> str:
+    # Convert to absolute path if it's not already absolute
+    # Handle both Unix-style absolute paths (starting with /) and Windows-style absolute paths
+    if not (notebook_path.startswith('/') or (len(notebook_path) > 1 and notebook_path[1] == ':')):
+        notebook_path = os.path.join(os.getcwd(), notebook_path)
+    return notebook_path
+
+def clean_directory_check(notebook_path: str) -> None:
+    notebook_path = resolve_notebook_path(notebook_path)
+    # pathlib handles the cross OS path conversion automatically
+    path = Path(notebook_path).resolve()
+    dir_path = path.parent
+
+    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

@@ -0,0 +1,116 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import sys
+import os
+import time
+import requests
+import tempfile
+import shutil
+import traceback
+import ast
+import importlib.util
+import warnings
+from typing import List, Tuple, Optional, Dict, Any, Generator
+from streamlit.testing.v1 import AppTest
+from contextlib import contextmanager
+from mito_ai.streamlit_conversion.streamlit_utils import resolve_notebook_path
+
+
+# warnings.filterwarnings("ignore", message=r".*missing ScriptRunContext.*")
+# warnings.filterwarnings("ignore", category=UserWarning)
+
+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(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)
+        
+        with change_working_directory(directory):
+            app_test = AppTest.from_string(app_code, 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
+
+    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})
+
+            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)})
+
+        finally:
+            self.cleanup()
+
+        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)
+    
+    has_validation_error = len(errors) > 0
+    stringified_errors = [str(error) for error in errors]
+    return has_validation_error, stringified_errors

mito_ai/streamlit_preview/handlers.py

@@ -70,8 +70,8 @@ class StreamlitPreviewHandler(APIHandler):
         
         # If still not found, return the original path (will cause a clear error)
         # This ensures we get a meaningful error message rather than a generic "file not found"
-        return notebook_path
-    
+        return os.path.join(os.getcwd(), notebook_path)
+
     @tornado.web.authenticated
     async def post(self) -> None:
         """Start a new streamlit preview.
@@ -97,7 +97,7 @@ class StreamlitPreviewHandler(APIHandler):
                 return
 
             notebook_path = body.get('notebook_path')
-            
+
             if not notebook_path:
                 self.set_status(400)
                 self.finish({"error": 'Missing notebook_path parameter'})
@@ -110,6 +110,7 @@ class StreamlitPreviewHandler(APIHandler):
             preview_id = str(uuid.uuid4())
             
             # Generate streamlit code using existing handler
+            print('notebook_path', notebook_path)
             success, app_path, message = await streamlit_handler(resolved_notebook_path)
             
             if not success or app_path is None:
@@ -136,7 +137,9 @@ class StreamlitPreviewHandler(APIHandler):
         except Exception as e:
             print(f"Error in streamlit preview handler: {e}")
             self.set_status(500)
-            self.finish({"error": f'Internal server error: {str(e)}'})
+            
+            # Respond with the error
+            self.finish({"error": str(e)})
     
     @tornado.web.authenticated
     def delete(self, preview_id: str) -> None: