Merge pull request #2147 from mito-ds/dev

Release Jan 14, 2026

Author: ngafar

.cursor/commands/verify.md

@@ -0,0 +1,58 @@
+# Verify Feature Command
+
+When invoked, automatically test and fix the feature that you just created using the Browser tool to access JupyterLab. Continue correcting and verifying the feature until it works. 
+
+## Execution Steps
+
+### 0. Setting up Environment
+
+Make sure that the code passes linting first. Check the typescript terminal that is running for errors. You must resolve those first to ensure you are testing the most up to date version of the code.
+
+### 1. Check JupyterLab Status
+
+First, verify JupyterLab is running on port 8888. Navigate to:
+```
+http://localhost:8888/lab?token=dev-token-for-cursor-testing-12345
+```
+
+### 2. Navigate and Load
+
+Use browser to navigate to JupyterLab URL above. Wait 3-5 seconds for full page load, then take a snapshot to understand current state.
+
+### 3. Identify Feature to Test
+
+Use the conversation history to identify the feature to test. Your goal is to QA the new feature that was just created.
+
+### 4. Test the Feature
+
+Interact with the feature using browser tools:
+
+### 5. Check for Issues
+
+After each interaction:
+- **Browser console**: Use `browser_console_messages` to check for JavaScript/TypeScript errors
+- **Network requests**: Use `browser_network_requests` to verify API calls succeed
+- **UI state**: Take snapshots to verify elements render correctly
+- **Functionality**: Verify feature works as intended
+
+### 6. Fix Issues Found
+
+- Fix the issues you found if if they are related to the feature you are testing
+- Document any additional issues you found that are not related to the feature you are testing
+
+### 7. Rebuild and Re-test
+
+After making fixes:
+
+1. **Setup Environment**: 
+   - For frontend changes: Wait a few seconds for rebuild to complete
+   - For backend changes: Restart and relaunch the server 
+
+2. **Refresh browser**: 
+   - Take new snapshot
+   - Re-test the feature
+   - Verify fix worked
+
+### 8. Iterate Until Complete
+
+Repeat steps 4-7 until the feature works correctly.

.gitignore

@@ -1,6 +1,3 @@
-# Ignore all the scratch files
-scratch*
-
 # Ignore all the ipynb files in mitosheet and mito-ai
 mitosheet/**/*.ipynb
 mito-ai/**/*.ipynb

mito-ai/mito_ai/__init__.py

@@ -19,6 +19,7 @@
 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/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/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/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/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/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/mito_ai/completions/handlers.py

@@ -34,6 +34,7 @@
     CodeExplainMetadata,
     AgentExecutionMetadata,
     InlineCompleterMetadata,
+    ScratchpadResultMetadata,
     MessageType
 )
 from mito_ai.completions.providers import OpenAIProvider
@@ -45,6 +46,7 @@
 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 @@ async def on_message(self, message: str) -> None: # type: ignore
             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/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:
     """