mito-ai 0.1.56__py3-none-any.whl → 0.1.57__py3-none-any.whl

mito_ai/__init__.py

@@ -19,6 +19,7 @@ from mito_ai.app_manager.handlers import AppManagerHandler
 from mito_ai.file_uploads.urls import get_file_uploads_urls
 from mito_ai.user.urls import get_user_urls
 from mito_ai.chat_history.urls import get_chat_history_urls
+from mito_ai.chart_wizard.urls import get_chart_wizard_urls
 
 # Force Matplotlib to use the Jupyter inline backend.
 # Background: importing Streamlit sets os.environ["MPLBACKEND"] = "Agg" very early.
@@ -109,6 +110,7 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     handlers.extend(get_file_uploads_urls(base_url)) # type: ignore
     handlers.extend(get_user_urls(base_url)) # type: ignore
     handlers.extend(get_chat_history_urls(base_url, global_message_history)) # type: ignore
+    handlers.extend(get_chart_wizard_urls(base_url, open_ai_provider)) # type: ignore
 
     web_app.add_handlers(host_pattern, handlers)
     server_app.log.info("Loaded the mito_ai server extension")

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.56'
+__version__ = VERSION = '0.1.57'

mito_ai/chart_wizard/__init__.py

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

mito_ai/chart_wizard/handlers.py

@@ -0,0 +1,52 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import tornado
+from typing import List
+from jupyter_server.base.handlers import APIHandler
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.utils.anthropic_utils import FAST_ANTHROPIC_MODEL
+from mito_ai.completions.models import MessageType
+from mito_ai.completions.prompt_builders.chart_conversion_prompt import create_chart_conversion_prompt
+
+class ChartWizardHandler(APIHandler):
+    def initialize(self, llm: OpenAIProvider) -> None:
+        """Initialize the handler with the LLM provider."""
+        super().initialize()
+        self._llm = llm
+
+    @tornado.web.authenticated
+    async def post(self) -> None:
+        """POST endpoint that receives code from the frontend and sends it to LLM."""
+        try:
+            data = json.loads(self.request.body.decode('utf-8'))
+            code = data.get('code', '')
+            
+            # Create prompt using the prompt builder
+            prompt = create_chart_conversion_prompt(code)
+            
+            # Call LLM
+            messages: List[ChatCompletionMessageParam] = [{"role": "user", "content": prompt}]
+            converted_code = await self._llm.request_completions(
+                messages=messages,
+                model=FAST_ANTHROPIC_MODEL,
+                message_type=MessageType.CHAT,
+                thread_id=None
+            )
+            
+            # Return the converted code
+            self.write({
+                "message": "Code converted successfully",
+                "converted_code": converted_code
+            })
+            self.finish()
+        except json.JSONDecodeError:
+            self.set_status(400)
+            self.write({"error": "Invalid JSON in request body"})
+            self.finish()
+        except Exception as e:
+            self.set_status(500)
+            self.write({"error": str(e)})
+            self.finish()

mito_ai/chart_wizard/urls.py

@@ -0,0 +1,23 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple, Any
+from jupyter_server.utils import url_path_join
+from mito_ai.chart_wizard.handlers import ChartWizardHandler
+from mito_ai.completions.providers import OpenAIProvider
+
+
+def get_chart_wizard_urls(base_url: str, llm: OpenAIProvider) -> List[Tuple[str, Any, dict]]:
+    """Get all chart wizard related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+        llm: The OpenAI provider instance
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    return [
+        (url_path_join(BASE_URL, "chart-wizard"), ChartWizardHandler, {"llm": llm}),
+    ]

mito_ai/completions/completion_handlers/completion_handler.py

@@ -3,11 +3,20 @@
 
 from typing import Protocol, TypeVar
 from abc import abstractmethod, ABCMeta
-from mito_ai.completions.models import ChatMessageMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, InlineCompleterMetadata, AgentSmartDebugMetadata
+from mito_ai.completions.models import ChatMessageMetadata, ScratchpadResultMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, InlineCompleterMetadata, AgentSmartDebugMetadata
 from mito_ai.completions.providers import OpenAIProvider
 from mito_ai.completions.message_history import GlobalMessageHistory
 
-T = TypeVar('T', ChatMessageMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, AgentSmartDebugMetadata, InlineCompleterMetadata, contravariant=True)
+T = TypeVar('T', 
+            ChatMessageMetadata, 
+            SmartDebugMetadata, 
+            CodeExplainMetadata, 
+            AgentExecutionMetadata, 
+            AgentSmartDebugMetadata, 
+            InlineCompleterMetadata,
+            ScratchpadResultMetadata,
+            contravariant=True
+        )
 
 class CompletionHandler(Protocol[T], metaclass=ABCMeta):
     """Protocol defining the interface for completion handlers.

mito_ai/completions/completion_handlers/scratchpad_result_handler.py

@@ -0,0 +1,66 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Literal, Union
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.models import ScratchpadResultMetadata, MessageType, ResponseFormatInfo, AgentResponse
+from mito_ai.completions.prompt_builders.scratchpad_result_prompt import create_scratchpad_result_prompt
+from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.completions.message_history import GlobalMessageHistory
+from mito_ai.completions.completion_handlers.completion_handler import CompletionHandler
+from mito_ai.completions.completion_handlers.utils import append_agent_system_message, create_ai_optimized_message
+
+__all__ = ["get_scratchpad_result_completion"]
+
+class ScratchpadResultHandler(CompletionHandler[ScratchpadResultMetadata]):
+    """Handler for scratchpad result completions."""
+    
+    @staticmethod
+    async def get_completion(
+        metadata: ScratchpadResultMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        model: str
+    ) -> str:
+        """Get a scratchpad result completion from the AI provider."""
+
+        if metadata.index is not None:
+            message_history.truncate_histories(
+                thread_id=metadata.threadId,
+                index=metadata.index
+            )
+
+        # Add the system message if it doesn't already exist
+        await append_agent_system_message(message_history, model, provider, metadata.threadId, True)
+        
+        # Create the prompt
+        prompt = create_scratchpad_result_prompt(metadata)
+        display_prompt = ""
+        
+        # Add the prompt to the message history
+        new_ai_optimized_message = create_ai_optimized_message(prompt, None, None)
+        new_display_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": display_prompt}
+        
+        await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, model, provider, metadata.threadId)
+        
+        # Get the completion
+        completion = await provider.request_completions(
+            messages=message_history.get_ai_optimized_history(metadata.threadId), 
+            model=model,
+            response_format_info=ResponseFormatInfo(
+                name='agent_response',
+                format=AgentResponse
+            ),
+            message_type=MessageType.AGENT_SCRATCHPAD_RESULT,
+            user_input="",
+            thread_id=metadata.threadId
+        )
+        
+        ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
+        
+        await message_history.append_message(ai_response_message, ai_response_message, model, provider, metadata.threadId)
+        
+        return completion
+
+# Use the static method directly
+get_scratchpad_result_completion = ScratchpadResultHandler.get_completion

mito_ai/completions/handlers.py

@@ -34,6 +34,7 @@ from mito_ai.completions.models import (
     CodeExplainMetadata,
     AgentExecutionMetadata,
     InlineCompleterMetadata,
+    ScratchpadResultMetadata,
     MessageType
 )
 from mito_ai.completions.providers import OpenAIProvider
@@ -45,6 +46,7 @@ from mito_ai.completions.completion_handlers.code_explain_handler import get_cod
 from mito_ai.completions.completion_handlers.inline_completer_handler import get_inline_completion
 from mito_ai.completions.completion_handlers.agent_execution_handler import get_agent_execution_completion
 from mito_ai.completions.completion_handlers.agent_auto_error_fixup_handler import get_agent_auto_error_fixup_completion
+from mito_ai.completions.completion_handlers.scratchpad_result_handler import get_scratchpad_result_completion
 from mito_ai.utils.telemetry_utils import identify
 
 FALLBACK_MODEL = "gpt-4.1"  # Default model to use for safety
@@ -314,6 +316,9 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             elif type == MessageType.AGENT_AUTO_ERROR_FIXUP:
                 agent_auto_error_fixup_metadata = AgentSmartDebugMetadata(**metadata_dict)
                 completion = await get_agent_auto_error_fixup_completion(agent_auto_error_fixup_metadata, self._llm, self._message_history, model)
+            elif type == MessageType.AGENT_SCRATCHPAD_RESULT:
+                scratchpad_result_metadata = ScratchpadResultMetadata(**metadata_dict)
+                completion = await get_scratchpad_result_completion(scratchpad_result_metadata, self._llm, self._message_history, model)
             elif type == MessageType.INLINE_COMPLETION:
                 inline_completer_metadata = InlineCompleterMetadata(**metadata_dict)
                 completion = await get_inline_completion(inline_completer_metadata, self._llm, self._message_history, model)

mito_ai/completions/models.py

@@ -29,13 +29,26 @@ class CellUpdate(BaseModel):
 # for now and rely on the AI to respond with the correct types, following the format
 # that we show it in the system prompt.
 class AgentResponse(BaseModel):
-    type: Literal['cell_update', 'get_cell_output', 'run_all_cells', 'finished_task', 'create_streamlit_app', 'edit_streamlit_app']
+    type: Literal[
+        'cell_update', 
+        'get_cell_output', 
+        'run_all_cells', 
+        'finished_task', 
+        'create_streamlit_app', 
+        'edit_streamlit_app', 
+        'ask_user_question', 
+        'scratchpad',
+    ]
     message: str
     cell_update: Optional[CellUpdate]
     get_cell_output_cell_id: Optional[str]
     next_steps: Optional[List[str]]
     analysis_assumptions: Optional[List[str]]
     streamlit_app_prompt: Optional[str]
+    question: Optional[str]
+    answers: Optional[List[str]]
+    scratchpad_code: Optional[str]
+    scratchpad_summary: Optional[str]
     
     
 @dataclass(frozen=True)
@@ -67,6 +80,7 @@ class MessageType(Enum):
     STREAMLIT_CONVERSION = "streamlit_conversion"
     STOP_AGENT = "stop_agent"
     DEPLOY_APP = "deploy_app"
+    AGENT_SCRATCHPAD_RESULT = "agent:scratchpad-result"
 
     
 @dataclass(frozen=True)
@@ -136,13 +150,20 @@ class CodeExplainMetadata():
     activeCellCode: Optional[str] = None
     
 @dataclass(frozen=True)
-class InlineCompleterMetadata():
+class InlineCompleterMetadata():    
     promptType: Literal['inline_completion']
     prefix: str 
     suffix: str
     variables: Optional[List[str]] = None
     files: Optional[List[str]] = None
 
+@dataclass(frozen=True)
+class ScratchpadResultMetadata():
+    promptType: Literal['agent:scratchpad-result']
+    threadId: ThreadID
+    scratchpadResult: str
+    index: Optional[int] = None
+
 @dataclass(frozen=True)
 class CompletionRequest:
     """

mito_ai/completions/prompt_builders/agent_smart_debug_prompt.py

@@ -22,12 +22,12 @@ def create_agent_smart_debug_prompt(md: AgentSmartDebugMetadata) -> str:
 
 Use this strategy for this message only. After this message, continue using the original set of instructions that I provided you.
 
-It is very important that When fixing this error, you do not change the original intent of the code cell. 
+It is very important that when fixing this error, you do not change the original intent of the code cell. 
 
 To fix this error, take the following approach: 
 Step 1: ERROR ANALYSIS: Analyze the error message to identify why the code cell errored.
 Step 2: INTENT PRESERVATION: Make sure you understand the intent of the CELL_UPDATE so that you can be sure to preserve it when you create a new CELL_UPDATE
-Step 3: ERROR CORRECTION: Respond with a new CELL_UPDATE that is applied to the same cell as the erroring CELL_UPDATE.
+Step 3: ERROR CORRECTION: Respond with a new CELL_UPDATE that is applied to the same cell as the erroring CELL_UPDATE or use the ASK_USER_QUESTION tool to get more information about how to proceed.
 
 INSTRUCTIONS FOR EACH PHASE
 
@@ -43,13 +43,15 @@ INTENT PRESERVATION:
 
 ERROR CORRECTION:
 
-- Return the full, updated version of cell {md.error_message_producing_code_cell_id} with the error fixed and a short explanation of the error.
+- Use one of your tools to correct the error or get more information from the user on how to proceed.
+- If you use the CELL_UPDATE tool, you must reutn the full updated version of cell {md.error_message_producing_code_cell_id} with the error fixed and a short explanation of the error.
 - You can only update code in {md.error_message_producing_code_cell_id}. You are unable to edit the code in any other cell when resolving this error.
 - Propose a solution that fixes the error and does not change the user's intent.
 - Make the solution as simple as possible.
 - Reuse as much of the existing code as possible.
 - DO NOT ADD TEMPORARY COMMENTS like '# Fixed the typo here' or '# Added this line to fix the error'
 - If you encounter a ModuleNotFoundError, you can install the package by adding the the following line to the top of the code cell: `!pip install <package_name> --quiet`.
+- If the error is not resolvable without getting more information from the user, you can respond with a ASK_USER_QUESTION tool call.
 - If you encounter a NameError, you can use the RUN_ALL_CELLS tool to run all cells from the top of the notebook to the bottom to bring the variable into scope.
     RUN_ALL_CELLS:
     When you want to execute all cells in the notebook from top to bottom, respond with this format:

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -4,13 +4,13 @@
 from typing import List
 from mito_ai.completions.prompt_builders.prompt_section_registry import SG, Prompt
 from mito_ai.completions.prompt_builders.prompt_constants import (
+    CHART_CONFIG_RULES,
     CITATION_RULES,
     CELL_REFERENCE_RULES,
     get_database_rules
 )
 from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
 
-
 def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
     
     # The GET_CELL_OUTPUT tool only works on Chrome based browsers. 
@@ -28,6 +28,8 @@ The user is going to ask you to guide them as they complete a task. You will hel
 You have access to a set of tools that you can use to accomplish the task you've been given. You can use one tool per message, and will receive the result of that tool use in the user's response. You use tools step-by-step to accomplish a given task, with each tool use informed by the result of the previous tool use.
 
 Each time you use a tool, except for the finished_task tool, the user will execute the tool and provide you with updated information about the notebook and variables defined in the kernel to help you decide what to do next."""))
+    
+    sections.append(SG.Generic("Chart Config Rules", CHART_CONFIG_RULES))
 
     sections.append(SG.Generic("TOOL: CELL_UPDATE", """
 
@@ -223,6 +225,85 @@ Important information:
 4. If running all cells results in an error, the system will automatically handle the error through the normal error fixing process.
 5. Do not use this tool repeatedly if it continues to produce errors - instead, focus on fixing the specific error that occurred."""))
     
+    # SCRATCHPAD tool
+    sections.append(SG.Generic("TOOL: SCRATCHPAD", """
+When you need to explore data, check the filesystem, analyze mappings, or look up values without leaving code in the notebook, use the SCRATCHPAD tool.
+
+Format:
+{{
+    type: 'scratchpad',
+    message: str,
+    scratchpad_code: str,
+    scratchpad_summary: str
+}}
+
+Important information:
+1. The scratchpad_code will execute silently against the same kernel as your notebook, so you have access to all variables and dataframes.
+2. Any variables you create in scratchpad code MUST be prefixed with "scratch_" (e.g., use "scratch_temp_df" not "temp_df", use "scratch_files" not "files"). This prevents them from appearing in the variable list and confusing future decisions.
+3. CRITICAL: Do NOT modify existing variables. If you need to explore or modify data, create a copy with the scratch_ prefix first. For example, use "scratch_df = df.copy()" and then modify scratch_df, rather than modifying df directly. This ensures existing variables remain unchanged.
+4. Structure your code to print the information you need. Use print() statements for output that will be captured.
+5. If you need structured data, consider using JSON: `import json; print(json.dumps(your_data))`
+6. The results (including any errors) will be included in your next message, so you can use them to inform your next action.
+7. If the code errors, the error message and traceback will be included in the results. You can then decide to fix the code and try again, ask the user a question, or take a different approach.
+8. Use scratchpad for exploration work that doesn't belong in the final notebook. Once you have the information, create clean CELL_UPDATES with hardcoded values.
+9. The scratchpad_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Checking files", "Exploring data", "Analyzing mappings", "Looking up values"). Avoid full sentences or explanations. This should read like a quick commit message or code label, not a description.
+
+Example:
+{{
+    type: 'scratchpad',
+    message: "I'll check what files are in the current directory to find the data file.",
+    scratchpad_code: "import os\\nscratch_files = os.listdir('.')\\nprint('Files:', scratch_files)\\nfor scratch_file in scratch_files:\\n    if scratch_file.endswith('.csv'):\\n        print(f'CSV file found: {scratch_file}')",
+    scratchpad_summary: "Checking files"
+}}
+
+Example with dataframe exploration:
+{{
+    type: 'scratchpad',
+    message: "I'll explore the dataframe structure to understand the columns.",
+    scratchpad_code: "scratch_df = df.copy()\\nprint('Columns:', scratch_df.columns.tolist())\\nprint('Shape:', scratch_df.shape)\\nprint('First few rows:')\\nprint(scratch_df.head())",
+    scratchpad_summary: "Exploring dataframe"
+}}
+"""))
+    
+    # ASK_USER_QUESTION tool 
+    sections.append(SG.Generic("TOOL: ASK_USER_QUESTION", f"""
+
+When you have a specific question that you the user to answer so that you can figure out how to proceed in your work, you can respond in this format:
+
+{{
+    type: 'ask_user_question',
+    message: str,
+    question: str,
+    answers: Optional[List[str]]
+}}
+
+Important information:
+1. Use this tool when you need clarification from the user on how to proceed. Common scenarios include:
+   - A file or resource doesn't exist and you need to know how to proceed
+   - There are multiple valid approaches and you want the user to choose
+   - You need clarification on ambiguous requirements
+   - You encounter an error that requires user input to resolve
+2. The message should be a short description of what you've tried so far and why you need to ask the user a question now. This helps the user understand the context. The message provides background information but should NOT contain the actual question.
+3. The question field is REQUIRED and must always be provided. It cannot be null or empty. The question should be a clear, direct question that ends with a question mark. It should be concise and direct - do NOT include instructions or explanations in the question itself, as the answer options should make it clear what information is needed. For example, instead of "Which companies would you like me to compare Meta's stock performance against? Please provide a list of company names or stock tickers", just ask "Which companies would you like me to compare Meta's stock performance against?" The answer options will make it clear that company names or tickers are expected.
+4. The message and question serve different purposes: the message provides context about what you've tried, while the question is the actual question the user needs to answer. If your message contains a question, extract that question into the question field. For example, if your message says "I need to understand how you'd like to access the tweets", your question should be something like "How would you like to access the tweets?"
+5. Use the optional list of answers to provide the user multiple choice options to choose from. If it is an open ended question that you cannot create helpful multiple choice answers for, leave it blank and the user will respond in the text input field. 
+6. When creating multiple choice answer options:
+   - Make each option distinct and meaningful - avoid options that differ only slightly from each other.
+   - If there are no obvious predefined answers, leave it blank and the user will respond in the text input field.
+7. After the user responds to your question, you will receive their response in the next message and can continue with the task based on their answer.
+8. Do not use this tool for trivial questions that you can infer from context. Only use it when you cannot proceed in the user's task without answering your specific question first.
+
+    <Example>
+    {{
+        type: 'ask_user_question',
+        message: "I tried importing apple_prices.csv and confirmed that it does not exist in the current working directory.",
+        question: "The file apple_prices.csv does not exist. How do you want to proceed?",
+        answers: ["Pull Apple Stock prices using yfinance API", "Create placeholder data", "Skip this step"]
+    }}
+    </Example>
+
+"""))
+    
     # CREATE_STREAMLIT_APP tool
     sections.append(SG.Generic("TOOL: CREATE_STREAMLIT_APP", """
 
@@ -316,18 +397,29 @@ Important information:
     # RULES section
     sections.append(SG.Generic("RULES", """
 - You are working in a Jupyter Lab environment in a .ipynb file.
-- In each message you can choose one of the tools to respond with. BUT YOU WILL GET TO SEND MULTIPLE MESSAGES TO THE USER TO ACCOMPLISH YOUR TASK SO DO NOT TRY TO ACCOMPLISH YOUR TASK IN A SINGLE MESSAGE.
+- In each message you can choose one of the tools to respond with. YOU WILL GET TO SEND MULTIPLE MESSAGES TO THE USER TO ACCOMPLISH YOUR TASK SO DO NOT TRY TO ACCOMPLISH YOUR TASK IN A SINGLE MESSAGE.
 - After you send a CELL_UPDATE, the user will send you a message with the updated variables, code, and files in the current directory. You will use this information to decide what to do next, so it is critical that you wait for the user's response after each CELL_UPDATE before deciding your next action.
-- When updating code, keep as much of the original code as possible and do not recreate variables that already exist.
 - When writing the message, do not explain to the user how to use the CELL_UPDATE or FINISHED_TASK response, they will already know how to use them. Just provide a summary of your thought process. Do not reference any Cell IDs in the message.
 - When writing the message, do not include leading words like "Explanation:" or "Thought process:". Just provide a summary of your thought process.
 - When writing the message, use tickmarks when referencing specific variable names. For example, write `sales_df` instead of "sales_df" or just sales_df."""))
     
     # CODE STYLE section
     sections.append(SG.Generic("CODE STYLE", """
-- Avoid using try/except blocks and other defensive programming patterns (like checking if files exist before reading them, verifying variables are defined before using them, etc.) unless there is a really good reason. In Jupyter notebooks, errors should surface immediately so users can identify and fix issues. When errors are caught and suppressed or when defensive checks hide problems, users continue running broken code without realizing it, and the agent's auto-error-fix loop cannot trigger. If a column doesn't exist, a file is missing, a variable isn't defined, or a module isn't installed, let it error. The user needs to know.
+- When updating code, keep as much of the original code as possible and do not recreate variables that already exist.
 - When you want to display a dataframe to the user, just write the dataframe on the last line of the code cell instead of writing print(<dataframe name>). Jupyter will automatically display the dataframe in the notebook.
-- When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter."""))
+- When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter.
+- Avoid adding try/except blocks unless there is a very good reason. Do not use them for things like: 
+    ```
+    try: 
+        df = pd.read_csv('my_data.csv')
+    except: 
+        print("File not found")
+    ```
+    Instead, just let the cell error and use the ask_user_question tool to figure out how to proceed.
+- Avoid defensive if statements like checking if a variable exists in the globals or verifying that a column exists. Instead, just let the code error and use the ask_user_question tool to figure out how to proceed.
+- Do not simulate the data without the user explicity asking you to do so.
+- Do not replace broken code with print statements that explain the issue. Instead, leave the broken code in the notebook and use the ask_user_question tools to communicate the issue to the user and figure out how to proceed.
+"""))
     
     # CITATION_RULES 
     sections.append(SG.Generic("Citation Rules", f"""{CITATION_RULES}

mito_ai/completions/prompt_builders/chart_conversion_prompt.py

@@ -0,0 +1,27 @@
+# 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.completions.prompt_builders.prompt_section_registry import SG, Prompt
+from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
+from mito_ai.completions.prompt_builders.prompt_constants import CHART_CONFIG_RULES
+
+def create_chart_conversion_prompt(code: str) -> str:
+    """
+    Create a prompt for converting matplotlib chart code to be used with the Chart Wizard.
+    
+    Args:
+        code: The matplotlib chart code to convert
+        
+    Returns:
+        A formatted prompt string
+    """
+    sections: List[PromptSection] = []
+
+    sections.append(SG.Generic("Instructions", "The following code contains a matplotlib chart. However, the chart must be converted to a specific format for use in our tool. Below you will find the rules used to create an acceptable chart; use these rules to reformat the code."))
+    
+    sections.append(SG.Generic("Chart Config Rules", CHART_CONFIG_RULES))
+    sections.append(SG.Generic("Code to Convert", f"```python\n{code}\n```"))
+    
+    prompt = Prompt(sections)
+    return str(prompt)

mito_ai/completions/prompt_builders/chat_system_message.py

@@ -4,6 +4,7 @@
 from typing import List
 from mito_ai.completions.prompt_builders.prompt_section_registry import SG, Prompt
 from mito_ai.completions.prompt_builders.prompt_constants import (
+    CHART_CONFIG_RULES,
     CHAT_CODE_FORMATTING_RULES,
     CITATION_RULES,
     CELL_REFERENCE_RULES,
@@ -31,6 +32,7 @@ Other useful information:
 2. If the user asks you to generate a dashboard, app, or streamlit app for them, you should tell them that they must use Agent mode to complete the task. You are not able to automatically switch the user to agent mode, but they can switch to it themselves by using the Chat/Agent mode toggle in the bottom left corner of the Ai taskpane.
 """))
 
+    sections.append(SG.Generic("Chart Config Rules", CHART_CONFIG_RULES))
     sections.append(SG.Generic("DatabaseRules", get_database_rules()))
     sections.append(SG.Generic("Citation Rules", CITATION_RULES))
     sections.append(SG.Generic("Cell Reference Rules", CELL_REFERENCE_RULES))

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -11,6 +11,34 @@ import json
 from typing import Final
 from mito_ai.utils.schema import MITO_FOLDER
 
+CHART_CONFIG_RULES = """
+When creating a matplotlib chart, you must use the `# === CHART CONFIG ===` and `# === END CONFIG ===` markers to indicate the start and end of the chart configuration section.
+
+The chart configuration section is a list of variables used to customize the chart. This includes the titles, labels, colors, and any variables that affect the chart's appearance.
+
+Rules:
+- All imports must appear at the top, before the chart configuration section.
+- Variables with multiple words should be underscore-separated.
+- All colors should be in hex format (e.g., "#3498db"). Use quotes around the hex string: COLOR = "#3498db" or COLOR = '#3498db'. Do NOT nest quotes.
+- Variables can only be strings, numbers, booleans, tuples, or lists. 
+- NEVER include comments on the same line as a variable assignment. Each variable assignment must be on its own line with no trailing comments.
+- For string values, use either single or double quotes (e.g., TITLE = "Sales by Product" or TITLE = 'Sales by Product'). Do not use nested quotes (e.g., do NOT use '"value"').
+
+Common Mistakes to Avoid:
+- WRONG: COLOR = '"#1877F2" # Meta Blue'  (nested quotes and inline comment)
+- WRONG: COLOR = "#1877F2" # Meta Blue  (inline comment)
+- WRONG: COLOR = '"#1877F2"'  (nested quotes)
+- CORRECT: COLOR = "#1877F2"  (simple hex string, no nested quotes, no inline comments)
+
+Example:
+# === CHART CONFIG ===
+TITLE = "Sales by Product"
+X_LABEL = "Product"
+Y_LABEL = "Sales"
+BAR_COLOR = "#000000"
+# === END CONFIG ===
+"""
+
 CITATION_RULES = """
 It is important that the user is able to verify any insights that you share with them about their data. To make this easy for the user, you must cite the lines of code that you are drawing the insight from. To provide a citation, use one of the following formats inline in your response:
 

mito_ai/completions/prompt_builders/scratchpad_result_prompt.py

@@ -0,0 +1,17 @@
+# 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.completions.models import ScratchpadResultMetadata
+from mito_ai.completions.prompt_builders.prompt_section_registry import SG, Prompt
+from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
+
+
+def create_scratchpad_result_prompt(md: ScratchpadResultMetadata) -> str:
+    sections: List[PromptSection] = [
+        SG.Generic("Reminder", "Continue working on your current task using the scratchpad results below."),
+        SG.Generic("Scratchpad Result", f"The result of your scratchpad is: {md.scratchpadResult}"),
+    ]
+
+    prompt = Prompt(sections)
+    return str(prompt)

mito_ai/utils/provider_utils.py

@@ -36,7 +36,14 @@ def does_message_require_fast_model(message_type: MessageType) -> bool:
     so they don't slow down the user's experience.
     """
     
-    if message_type in (MessageType.CHAT, MessageType.SMART_DEBUG, MessageType.CODE_EXPLAIN, MessageType.AGENT_EXECUTION, MessageType.AGENT_AUTO_ERROR_FIXUP):
+    if message_type in (
+        MessageType.CHAT,
+        MessageType.SMART_DEBUG,
+        MessageType.CODE_EXPLAIN,
+        MessageType.AGENT_EXECUTION,
+        MessageType.AGENT_SCRATCHPAD_RESULT,
+        MessageType.AGENT_AUTO_ERROR_FIXUP,
+    ):
         return False
     elif message_type in (MessageType.INLINE_COMPLETION, MessageType.CHAT_NAME_GENERATION):
         return True