mito-ai 0.1.50__py3-none-any.whl

mito_ai/streamlit_conversion/agent_utils.py

@@ -0,0 +1,37 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple
+import re
+from anthropic.types import MessageParam
+from mito_ai.streamlit_conversion.prompts.streamlit_system_prompt import streamlit_system_prompt
+from mito_ai.utils.anthropic_utils import stream_anthropic_completion_from_mito_server
+from mito_ai.streamlit_conversion.prompts.prompt_constants import MITO_TODO_PLACEHOLDER
+from mito_ai.completions.models import MessageType
+
+STREAMLIT_AI_MODEL = "claude-sonnet-4-5-20250929"
+
+def extract_todo_placeholders(agent_response: str) -> List[str]:
+    """Extract TODO placeholders from the agent's response"""
+    return [line.strip() for line in agent_response.split('\n') if MITO_TODO_PLACEHOLDER in line]
+
+async def get_response_from_agent(message_to_agent: List[MessageParam]) -> str:
+    """Gets the streaming response from the agent using the mito server"""
+    model = STREAMLIT_AI_MODEL
+    max_tokens = 64000 # TODO: If we move to haiku, we must reset this to 8192
+    temperature = 0.2
+
+    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

mito_ai/streamlit_conversion/prompts/prompt_constants.py

@@ -0,0 +1,172 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+MITO_TODO_PLACEHOLDER = "# MITO_TODO_PLACEHOLDER"
+
+search_replace_instructions = f"""
+RESPONSE FORMAT: You can edit the existing code using the **SEARCH_REPLACE format** for exact string matching and replacement.
+
+**STRUCTURE:**
+```search_replace
+>>>>>>> SEARCH
+[exact code currently in the file]
+=======
+[new code to replace it with]
+<<<<<<< REPLACE
+```
+
+**COMPONENTS:**
+```search_replace - This is the start of the search/replace block
+- `>>>>>>> SEARCH` - Exact text that EXISTS NOW in the file (7 chevrons)
+- `=======` - Separator between the search and replace blocks (7 equals signs)  
+- `<<<<<<< REPLACE` - Replacement text (7 chevrons)
+
+---
+
+**CRITICAL RULES - READ CAREFULLY:**
+
+1. **SEARCH = CURRENT STATE ONLY**
+   - The SEARCH block must contain ONLY code that currently exists in the file
+   - NEVER include new code, future code, or code you wish existed in the SEARCH block
+   - Copy exact text from the current file, character-for-character
+   
+2. **EXACT MATCHING REQUIRED**
+   - Every space, tab, newline must match perfectly
+   - Preserve exact indentation (spaces vs tabs)
+   - Include trailing newlines if present
+   - No approximations - even one character difference will fail
+
+3. **SIZE LIMITS**
+   - There are no size limits to each search/replace block, however, it is generally preferable to keep the SEARCH blocks small and focused on one change. 
+   - For large changes, use multiple smaller search/replace blocks
+
+4. **UNIQUENESS**
+   - Include enough context to make the SEARCH block unique
+   - If text appears multiple times, add surrounding lines
+   - Ensure there's only ONE match in the file
+
+5. **VERIFICATION CHECKLIST** (before generating each block):
+   ✓ Is every line in my SEARCH block currently in the file?
+   ✓ Did I copy the exact spacing and whitespace?
+   ✓ Will this match exactly once?
+   
+6. **SEARCH REPLACE BLOCK STRUCTURE**
+   - You must adhere to to the exact search_replace structure as shown in the examples.
+
+---
+
+**MULTIPLE REPLACEMENTS:**
+- You can include multiple search/replace blocks in one response
+- Each block is independent and processed separately
+- Use separate ```search_replace blocks for each change
+
+<Example 1: Updating existing content>
+
+```search_replace
+>>>>>>> SEARCH
+st.title("Old Title")
+=======
+st.title("New Title")
+<<<<<<< REPLACE
+```
+</Example 1>
+
+<Example 2: Adding new content>
+
+```search_replace
+>>>>>>> SEARCH
+st.title("My App")
+=======
+st.title("My App")
+st.header("Welcome")
+st.write("This is a test app")
+<<<<<<< REPLACE
+```
+</Example 2>
+
+<Example 3: Deleting existing content>
+
+```search_replace
+>>>>>>> SEARCH
+st.write("Old message")
+=======
+<<<<<<< REPLACE
+```
+</Example 3>
+
+<Example 4: Multiple replacements in one response>
+
+```search_replace
+>>>>>>> SEARCH
+st.title("Old Title")
+=======
+st.title("New Title")
+<<<<<<< REPLACE
+```
+
+```search_replace
+>>>>>>> SEARCH
+st.write("Old message")
+=======
+st.write("New message")
+<<<<<<< REPLACE
+```
+</Example 4>
+
+<Example 5: Using extra context to identify the correct code to replace>
+
+In the below example, assume that the code st.write("Old message") appears multiple times in the file, so we use extra context lines to identify the correct code to replace.
+
+```search_replace
+>>>>>>> SEARCH
+# This is a unique comment
+st.write("Old message")
+=======
+# This is a unique comment
+st.write("New message")
+<<<<<<< REPLACE
+```
+</Example 5>
+
+<Example 6: Search/replace while respecting whitespace and indentation>
+
+```search_replace
+>>>>>>> SEARCH
+data_list = [
+    {{'id': 1, 'name': 'Item A'}},
+    {MITO_TODO_PLACEHOLDER}: Add remaining entries from notebook
+]
+=======
+data_list = [
+    {{'id': 1, 'name': 'Item A'}},
+    {{'id': 2, 'name': 'Item B'}},
+    {{'id': 3, 'name': 'Item C'}},
+    {{'id': 4, 'name': 'Item D'}}
+]
+<<<<<<< REPLACE
+```
+</Example 6>
+
+<Example 7: Tab structure changes>
+
+```search_replace
+>>>>>>> SEARCH
+tab1, tab2 = st.tabs(["Cat", "Dog"])
+
+with tab1:
+    st.header("A cat")
+    st.image("https://static.streamlit.io/examples/cat.jpg", width=200)
+with tab2:
+    st.header("A dog")
+    st.image("https://static.streamlit.io/examples/dog.jpg", width=200)
+=======
+st.header("A cat")
+st.image("https://static.streamlit.io/examples/cat.jpg", width=200)
+st.header("A dog")
+st.image("https://static.streamlit.io/examples/dog.jpg", width=200)
+<<<<<<< REPLACE
+```
+</Example 7>
+
+Your response must consist **only** of valid search_replace blocks.
+"""

mito_ai/streamlit_conversion/prompts/prompt_utils.py

@@ -0,0 +1,10 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+def add_line_numbers_to_code(code: str) -> str:
+    """Add line numbers to the code"""
+    code_with_line_numbers = ""
+    for i, line in enumerate(code.split('\n'), 1):
+        code_with_line_numbers += f"{i:3d}: {line}\n"
+        
+    return code_with_line_numbers

mito_ai/streamlit_conversion/prompts/streamlit_app_creation_prompt.py

@@ -0,0 +1,46 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List
+from mito_ai.streamlit_conversion.prompts.prompt_constants import MITO_TODO_PLACEHOLDER
+
+def get_streamlit_app_creation_prompt(notebook: List[dict]) -> str:
+    """
+    This prompt is used to create a streamlit app from a notebook.
+    """
+    return f"""Convert the following Jupyter notebook into a Streamlit application.
+
+GOAL: Create a complete, runnable Streamlit app that accurately represents the notebook. It must completely convert the notebook. 
+
+TODO PLACEHOLDER RULES:
+If you decide to leave any TODOs, you must mark them with {MITO_TODO_PLACEHOLDER}. You should use {MITO_TODO_PLACEHOLDER} instead of comments like the following: 
+- # ... (include all mappings from the notebook)
+- # ... (include all violation codes from the notebook)
+- # Fill in the rest of the code here
+- # TODO: Add more code here
+- # TODO: Add the visualization code here
+
+For each TODO, use this exact format:
+{MITO_TODO_PLACEHOLDER}: <specific description of what needs to be added>
+
+IMPORTANT:
+- The app must still be RUNNABLE even with placeholders
+- Include enough sample data to show the structure
+- Do NOT use placeholders for small/medium content - include it directly
+- Do NOT use placeholders for file paths, imports, or core logic
+- Only use placeholders when absolutely necessary. Add all of the content directly as much as possible.
+
+<Example>
+If the notebook has a list of dictionaries with 50 entries, you would write:
+
+data = [
+    {{'id': 1, 'name': 'Item A', 'category': 'Type 1', 'value': 100}},
+    {{'id': 2, 'name': 'Item B', 'category': 'Type 2', 'value': 200}},
+    {MITO_TODO_PLACEHOLDER}: Add remaining entries from the data list
+]
+</Example>
+
+Notebook to convert:
+
+{notebook}
+"""

mito_ai/streamlit_conversion/prompts/streamlit_error_correction_prompt.py

@@ -0,0 +1,28 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.streamlit_conversion.prompts.prompt_constants import search_replace_instructions
+from mito_ai.streamlit_conversion.prompts.prompt_utils import add_line_numbers_to_code
+
+def get_streamlit_error_correction_prompt(error: str, streamlit_app_code: str) -> str:
+    
+    existing_streamlit_app_code_with_line_numbers = add_line_numbers_to_code(streamlit_app_code)
+    
+    return f"""You've created a Streamlit app, but it has an error in it when you try to run it.
+
+Your job is to fix the error now. Only fix the specific error that you are instructed to fix now. Do not fix other error that that you anticipate. You will be asked to fix other errors later.
+
+{search_replace_instructions}
+
+===============================================
+
+EXISTING STREAMLIT APP:
+{existing_streamlit_app_code_with_line_numbers}
+
+===============================================
+
+Please create a search/replace block that corrects this error. Please keep your fix concise:
+{error}
+
+"""
+

mito_ai/streamlit_conversion/prompts/streamlit_finish_todo_prompt.py

@@ -0,0 +1,45 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List
+from mito_ai.streamlit_conversion.prompts.prompt_constants import MITO_TODO_PLACEHOLDER, search_replace_instructions
+from mito_ai.streamlit_conversion.prompts.prompt_utils import add_line_numbers_to_code
+
+def get_finish_todo_prompt(notebook: List[dict], existing_streamlit_app_code: str, todo_placeholder: str) -> str:
+  
+    existing_streamlit_app_code_with_line_numbers = add_line_numbers_to_code(existing_streamlit_app_code)
+    
+    return f"""You've already created the first draft of a Streamlit app representation of a Jupyter notebook, but you left yourself some TODOs marked as `{MITO_TODO_PLACEHOLDER}`. 
+
+**CRITICAL COMPLETION REQUIREMENT:**
+You have ONE and ONLY ONE opportunity to complete this TODO. If you do not finish the entire task completely, the application will be broken and unusable. This is your final chance to get it right.
+
+**COMPLETION RULES:**
+1. **NEVER leave partial work** - If the TODO asks for a list with 100 items, provide ALL 100 items
+2. **NEVER use placeholders** - This is your only opportunity to fulfill this TODO, so do not leave yourself another TODO.
+3. **NEVER assume "good enough"** - Complete the task to 100% satisfaction
+4. **If the task seems large, that's exactly why it needs to be done now** - This is your only chance
+
+**HOW TO DETERMINE IF TASK IS COMPLETE:**
+- If building a list/dictionary: Include ALL items that should be in the final data structure
+- If creating functions: Implement ALL required functionality
+- If converting a visualization: Copy over ALL of the visualization code from the notebook, including all styling and formatting.
+
+{search_replace_instructions}
+
+===============================================
+
+Input Notebook that you are converting into the Streamlit app:
+{notebook}
+
+===============================================
+
+EXISTING STREAMLIT APP:
+{existing_streamlit_app_code_with_line_numbers}
+
+===============================================
+
+Please make the changes for this TODO. Only focus on this one TODO right now. You will be asked to fix others later:
+{todo_placeholder}
+
+"""

mito_ai/streamlit_conversion/prompts/streamlit_system_prompt.py

@@ -0,0 +1,56 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+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
+- Experienced in creating executive-ready dashboards for business stakeholders
+- Skilled in translating technical analysis into clear, interactive presentations
+
+TASK REQUIREMENTS:
+1. Convert Jupyter notebook content into a complete Streamlit application (app.py)
+2. Preserve ALL outputs from code cells and markdown cells as they appear in the notebook
+3. Maintain the logical flow and structure of the original analysis
+4. Create an executive-friendly dashboard suitable for company leadership
+
+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
+- 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
+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(\"\"\"
+    <style>
+        #MainMenu {visibility: hidden;}
+        .stAppDeployButton {display:none;}
+        footer {visibility: hidden;}
+        .stMainBlockContainer {padding: 2rem 1rem 2rem 1rem;}
+    </style>
+\"\"\", unsafe_allow_html=True)
+```
+
+OUTPUT FORMAT:
+- Output the complete, runnable app.py file.
+- Do not output any extra text, just give the python code. 
+
+"""

mito_ai/streamlit_conversion/prompts/update_existing_app_prompt.py

@@ -0,0 +1,50 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List
+from mito_ai.streamlit_conversion.prompts.prompt_constants import search_replace_instructions
+from mito_ai.streamlit_conversion.prompts.prompt_utils import add_line_numbers_to_code
+
+def get_update_existing_app_prompt(notebook: List[dict], streamlit_app_code: str, edit_prompt: str) -> str:
+    """
+    This prompt is used to update an existing streamlit app.
+    """
+    
+    existing_streamlit_app_code_with_line_numbers = add_line_numbers_to_code(streamlit_app_code)
+    
+    return f"""
+
+GOAL: You've previously created a first draft of the Streamlit app. Now the user reviewed it and provided feedback.Update the existing streamlit app according to the feedback provided by the user. Use the input notebook to help you understand what code needs to be added, changed, or modified to fulfill the user's edit request.
+
+**CRITICAL COMPLETION REQUIREMENT:**
+You have ONE and ONLY ONE opportunity to complete this edit request. If you do not finish the entire task completely, the application will be broken and unusable. This is your final chance to get it right.
+
+**COMPLETION RULES:**
+1. **NEVER leave partial work** - If the edit request requires generating a list with 100 items, provide ALL 100 items.
+2. **NEVER use placeholders** - This is your only opportunity to fulfill this edit request, so do not leave yourself another TODOs.
+3. **NEVER assume "good enough"** - Complete the task to 100% satisfaction.
+4. **If the task seems large, that's exactly why it needs to be done now** - This is your only chance
+
+**HOW TO DETERMINE IF TASK IS COMPLETE:**
+- If building a list/dictionary: Include ALL items that should be in the final data structure.
+- If creating functions: Implement ALL required functionality.
+- If converting a visualization: Copy over ALL of the visualization code from the notebook, including all styling and formatting.
+
+{search_replace_instructions}
+
+===============================================
+
+INPUT NOTEBOOK: 
+{notebook}
+
+===============================================
+
+EXISTING STREAMLIT APP:
+{existing_streamlit_app_code_with_line_numbers}
+
+===============================================
+
+USER EDIT REQUEST:
+{edit_prompt}
+
+"""

mito_ai/streamlit_conversion/search_replace_utils.py

@@ -0,0 +1,94 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import re
+from typing import List, Tuple
+
+from mito_ai.utils.error_classes import StreamlitConversionError
+from mito_ai.utils.telemetry_utils import log
+
+
+def extract_search_replace_blocks(message_content: str) -> List[Tuple[str, str]]:
+    """
+    Extract all search_replace blocks from Claude's response.
+    
+    Returns:
+        List of tuples (search_text, replace_text) for each search/replace block
+    """
+    if "```search_replace" not in message_content:
+        return []
+    
+    pattern = r'```search_replace\n(.*?)```'
+    matches = re.findall(pattern, message_content, re.DOTALL)
+    
+    search_replace_pairs = []
+    for match in matches:
+        # Split by the separator
+        if "=======" not in match:
+            continue
+            
+        parts = match.split("=======", 1)
+        if len(parts) != 2:
+            continue
+            
+        search_part = parts[0]
+        replace_part = parts[1]
+        
+        # Extract search text (after SEARCH marker)
+        if ">>>>>>> SEARCH" in search_part:
+            search_text = search_part.split(">>>>>>> SEARCH", 1)[1].strip()
+        else:
+            continue
+            
+        # Extract replace text (before REPLACE marker)
+        if "<<<<<<< REPLACE" in replace_part:
+            replace_text = replace_part.split("<<<<<<< REPLACE", 1)[0].strip()
+        else:
+            continue
+            
+        search_replace_pairs.append((search_text, replace_text))
+    
+    return search_replace_pairs
+
+
+def apply_search_replace(text: str, search_replace_pairs: List[Tuple[str, str]]) -> str:
+    """
+    Apply search/replace operations to the given text.
+    
+    Parameters
+    ----------
+    text : str
+        The original file contents.
+    search_replace_pairs : List[Tuple[str, str]]
+        List of (search_text, replace_text) tuples to apply.
+        
+    Returns
+    -------
+    str
+        The updated contents after applying all search/replace operations.
+        
+    Raises
+    ------
+    ValueError
+        If a search text is not found or found multiple times.
+    """
+    if not search_replace_pairs:
+        return text
+    
+    result = text
+    
+    for search_text, replace_text in search_replace_pairs:
+        # Count occurrences of search text
+        count = result.count(search_text)
+        
+        if count == 0:
+            print("Search Text Not Found: ", repr(search_text))
+            raise StreamlitConversionError(f"Search text not found: {repr(search_text)}", error_code=500)
+        elif count > 1:
+            print("Search Text Found Multiple Times: ", repr(search_text))
+            log("mito_ai_search_text_found_multiple_times", params={"search_text": repr(search_text)}, key_type="mito_server_key")
+
+        # Perform the replacement
+        result = result.replace(search_text, replace_text, 1)
+    
+    return result