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

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.providers import OpenAIProvider
+from mito_ai.completions.models import ChatMessageMetadata, ScratchpadResultMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, InlineCompleterMetadata, AgentSmartDebugMetadata
+from mito_ai.provider_manager import ProviderManager
 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.
@@ -20,9 +29,8 @@ class CompletionHandler(Protocol[T], metaclass=ABCMeta):
     @abstractmethod
     async def get_completion(
         metadata: T,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get a completion from the AI provider.
         
@@ -30,7 +38,6 @@ class CompletionHandler(Protocol[T], metaclass=ABCMeta):
             metadata: Metadata about the completion request
             provider: The AI provider to use
             message_history: The history of messages in the conversation
-            model: The model to use for the completion
             
         Returns:
             The completion string from the AI

mito_ai/completions/completion_handlers/inline_completer_handler.py

@@ -5,7 +5,7 @@ from typing import List
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import InlineCompleterMetadata, MessageType
 from mito_ai.completions.prompt_builders.inline_completer_prompt import create_inline_prompt
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.completions.completion_handlers.completion_handler import CompletionHandler
 
@@ -17,9 +17,8 @@ class InlineCompleterHandler(CompletionHandler[InlineCompleterMetadata]):
     @staticmethod
     async def get_completion(
         metadata: InlineCompleterMetadata,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get an inline completion from the AI provider."""
         
@@ -37,9 +36,9 @@ class InlineCompleterHandler(CompletionHandler[InlineCompleterMetadata]):
         # Get the completion
         completion = await provider.request_completions(
             messages=messages, 
-            model=model,
             message_type=MessageType.INLINE_COMPLETION,
-            thread_id=None
+            thread_id=None,
+            use_fast_model=True
         )
         
         return completion

mito_ai/completions/completion_handlers/scratchpad_result_handler.py

@@ -0,0 +1,64 @@
+# 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.provider_manager import ProviderManager
+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: ProviderManager,
+        message_history: GlobalMessageHistory
+    ) -> 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, 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, provider, metadata.threadId)
+        
+        # Get the completion
+        completion = await provider.request_completions(
+            messages=message_history.get_ai_optimized_history(metadata.threadId), 
+            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, provider, metadata.threadId)
+        
+        return completion
+
+# Use the static method directly
+get_scratchpad_result_completion = ScratchpadResultHandler.get_completion

mito_ai/completions/completion_handlers/smart_debug_handler.py

@@ -15,7 +15,7 @@ from mito_ai.completions.prompt_builders.smart_debug_prompt import (
     create_error_prompt,
     remove_inner_thoughts_from_message,
 )
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 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_chat_system_message
@@ -29,9 +29,8 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
     @staticmethod
     async def get_completion(
         metadata: SmartDebugMetadata,
-        provider: OpenAIProvider,
-        message_history: GlobalMessageHistory,
-        model: str
+        provider: ProviderManager,
+        message_history: GlobalMessageHistory
     ) -> str:
         """Get a smart debug completion from the AI provider."""
 
@@ -43,7 +42,7 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
         thread_id = metadata.threadId
 
         # Add the system message if it doesn't already exist
-        await append_chat_system_message(message_history, model, provider, thread_id)
+        await append_chat_system_message(message_history, provider, thread_id)
 
         # Create the prompt
         prompt = create_error_prompt(error_message, active_cell_code, active_cell_id, variables, files)
@@ -53,13 +52,12 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
         new_ai_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": prompt}
         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, thread_id
+            new_ai_optimized_message, new_display_optimized_message, provider, thread_id
         )
 
         # Get the completion
         completion = await provider.request_completions(
             messages=message_history.get_ai_optimized_history(thread_id),
-            model=model,
             message_type=MessageType.SMART_DEBUG,
             user_input=error_message,
             thread_id=thread_id
@@ -78,7 +76,7 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
             "content": display_completion,
         }
         await message_history.append_message(
-            ai_response_message, display_response_message, model, provider, thread_id
+            ai_response_message, display_response_message, provider, thread_id
         )
 
         return display_completion
@@ -86,11 +84,10 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
     @staticmethod
     async def stream_completion(
         metadata: SmartDebugMetadata,
-        provider: OpenAIProvider,
+        provider: ProviderManager,
         message_history: GlobalMessageHistory,
         message_id: str,
-        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
-        model: str
+        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]
     ) -> str:
         """Stream smart debug completions from the AI provider.
 
@@ -112,7 +109,7 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
         thread_id = metadata.threadId
 
         # Add the system message if it doesn't already exist
-        await append_chat_system_message(message_history, model, provider, thread_id)
+        await append_chat_system_message(message_history, provider, thread_id)
 
         # Create the prompt
         prompt = create_error_prompt(error_message, active_cell_code, active_cell_id, variables, files)
@@ -122,18 +119,17 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
         new_ai_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": prompt}
         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, thread_id
+            new_ai_optimized_message, new_display_optimized_message, provider, thread_id
         )
 
         # Stream the completions using the provider's stream method
         accumulated_response = await provider.stream_completions(
             message_type=MessageType.SMART_DEBUG,
             messages=message_history.get_ai_optimized_history(thread_id),
-            model=model,
             message_id=message_id,
+            thread_id=thread_id,
             reply_fn=reply_fn,
-            user_input=error_message,
-            thread_id=thread_id
+            user_input=error_message
         )
 
         # Process the completion to remove inner thoughts
@@ -149,7 +145,7 @@ class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
             "content": display_completion,
         }
         await message_history.append_message(
-            ai_response_message, display_response_message, model, provider, thread_id
+            ai_response_message, display_response_message, provider, thread_id
         )
 
         return display_completion

mito_ai/completions/completion_handlers/utils.py

@@ -5,7 +5,7 @@ import base64
 from typing import Optional, Union, List, Dict, Any, cast
 from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.completions.models import ThreadID
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.prompt_builders.chat_system_message import (
     create_chat_system_message_prompt,
@@ -17,8 +17,7 @@ from mito_ai.completions.prompt_builders.agent_system_message import (
 
 async def append_chat_system_message(
     message_history: GlobalMessageHistory,
-    model: str,
-    provider: OpenAIProvider,
+    provider: ProviderManager,
     thread_id: ThreadID,
 ) -> None:
 
@@ -39,7 +38,6 @@ async def append_chat_system_message(
     await message_history.append_message(
         ai_optimized_message=system_message,
         display_message=system_message,
-        model=model,
         llm_provider=provider,
         thread_id=thread_id,
     )
@@ -47,8 +45,7 @@ async def append_chat_system_message(
 
 async def append_agent_system_message(
     message_history: GlobalMessageHistory,
-    model: str,
-    provider: OpenAIProvider,
+    provider: ProviderManager,
     thread_id: ThreadID,
     isChromeBrowser: bool,
 ) -> None:
@@ -70,7 +67,6 @@ async def append_agent_system_message(
     await message_history.append_message(
         ai_optimized_message=system_message,
         display_message=system_message,
-        model=model,
         llm_provider=provider,
         thread_id=thread_id,
     )

mito_ai/completions/handlers.py

@@ -34,21 +34,22 @@ from mito_ai.completions.models import (
     CodeExplainMetadata,
     AgentExecutionMetadata,
     InlineCompleterMetadata,
+    ScratchpadResultMetadata,
     MessageType
 )
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.utils.create import initialize_user
 from mito_ai.utils.version_utils import is_pro
+from mito_ai.utils.model_utils import get_available_models
 from mito_ai.completions.completion_handlers.chat_completion_handler import get_chat_completion, stream_chat_completion
 from mito_ai.completions.completion_handlers.smart_debug_handler import get_smart_debug_completion, stream_smart_debug_completion
 from mito_ai.completions.completion_handlers.code_explain_handler import get_code_explain_completion, stream_code_explain_completion
 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
-
 # The GlobalMessageHistory is now created in __init__.py and passed to handlers
 # to ensure there's only one instance managing the .mito/ai-chats directory locks
 
@@ -59,13 +60,12 @@ FALLBACK_MODEL = "gpt-4.1"  # Default model to use for safety
 class CompletionHandler(JupyterHandler, WebSocketHandler):
     """Completion websocket handler."""
 
-    def initialize(self, llm: OpenAIProvider, message_history: GlobalMessageHistory) -> None:
+    def initialize(self, llm: ProviderManager, message_history: GlobalMessageHistory) -> None:
         super().initialize()
         self.log.debug("Initializing websocket connection %s", self.request.path)
         self._llm = llm
         self._message_history = message_history
         self.is_pro = is_pro()
-        self._selected_model = FALLBACK_MODEL
         self.is_electron = False
         identify(llm.key_type)
         
@@ -200,7 +200,25 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
         if type == MessageType.UPDATE_MODEL_CONFIG:
             model = metadata_dict.get('model')
             if model:
-                self._selected_model = model
+                # Validate model is in allowed list
+                available_models = get_available_models()
+                if model not in available_models:
+                    error = CompletionError(
+                        error_type="InvalidModelConfig",
+                        title="Invalid model configuration",
+                        traceback="",
+                        hint=f"Model '{model}' is not in the allowed model list. Available models: {', '.join(available_models)}"
+                    )
+                    reply = CompletionReply(
+                        items=[],
+                        error=error,
+                        parent_id=parsed_message.get('message_id')
+                    )
+                    self.reply(reply)
+                    return
+                
+                # Set the model in ProviderManager
+                self._llm.set_selected_model(model)
                 self.log.info(f"Model updated to: {model}")
                 reply = CompletionReply(
                     items=[CompletionItem(content=f"Model updated to {model}", isIncomplete=False)],
@@ -239,7 +257,6 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                 await self._message_history.append_message(
                     ai_optimized_message=ai_optimized_message,
                     display_message=display_optimized_message,
-                    model=self._selected_model,
                     llm_provider=self._llm,
                     thread_id=thread_id_to_stop
                 )
@@ -253,8 +270,6 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             message_id = parsed_message.get('message_id')
             stream = parsed_message.get('stream')
 
-            # When handling completions, always use the selected model
-            model = self._selected_model
             if type == MessageType.CHAT:
                 chat_metadata = ChatMessageMetadata(**metadata_dict)
                 
@@ -266,13 +281,12 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                         self._llm, 
                         self._message_history, 
                         message_id,
-                        self.reply,
-                        model
+                        self.reply
                     )
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_chat_completion(chat_metadata, self._llm, self._message_history, model)
+                    completion = await get_chat_completion(chat_metadata, self._llm, self._message_history)
             elif type == MessageType.SMART_DEBUG:
                 smart_debug_metadata = SmartDebugMetadata(**metadata_dict)
                 # Handle streaming if requested and available
@@ -283,13 +297,12 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                         self._llm, 
                         self._message_history, 
                         message_id,
-                        self.reply,
-                        model
+                        self.reply
                     )
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, self._message_history, model)
+                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, self._message_history)
             elif type == MessageType.CODE_EXPLAIN:
                 code_explain_metadata = CodeExplainMetadata(**metadata_dict)
 
@@ -301,22 +314,24 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                         self._llm, 
                         self._message_history,
                         message_id,
-                        self.reply,
-                        model
+                        self.reply
                     )
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, self._message_history, model)
+                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, self._message_history)
             elif type == MessageType.AGENT_EXECUTION:
                 agent_execution_metadata = AgentExecutionMetadata(**metadata_dict)
-                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, self._message_history, model)
+                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, self._message_history)
             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)
+                completion = await get_agent_auto_error_fixup_completion(agent_auto_error_fixup_metadata, self._llm, self._message_history)
+            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)
             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)
+                completion = await get_inline_completion(inline_completer_metadata, self._llm, self._message_history)
             else:
                 raise ValueError(f"Invalid message type: {type}")
             

mito_ai/completions/message_history.py

@@ -11,7 +11,7 @@ from typing import Dict, List, Optional
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.models import CompletionRequest, ChatThreadMetadata, MessageType, ThreadID
 from mito_ai.completions.prompt_builders.chat_name_prompt import create_chat_name_prompt
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.utils.schema import MITO_FOLDER
 from mito_ai.utils.message_history_utils import trim_old_messages
 
@@ -19,16 +19,15 @@ CHAT_HISTORY_VERSION = 2 # Increment this if the schema changes
 NEW_CHAT_NAME = "(New Chat)"
 NUMBER_OF_THREADS_CUT_OFF = 50
 
-async def generate_short_chat_name(user_message: str, assistant_message: str, model: str, llm_provider: OpenAIProvider) -> str:
+async def generate_short_chat_name(user_message: str, assistant_message: str, llm_provider: ProviderManager) -> str:
     prompt = create_chat_name_prompt(user_message, assistant_message)
 
     completion = await llm_provider.request_completions(
         messages=[{"role": "user", "content": prompt}], 
-        # We set the model so we can use the correct model provider, but request_completions will decide to 
-        # use the fast model from that provider to make the request.
-        model=model, 
+        # Use fast model from the selected provider for chat name generation
         message_type=MessageType.CHAT_NAME_GENERATION,
-        thread_id=None
+        thread_id=None,
+        use_fast_model=True
     )
     
     # Do a little cleanup of the completion. Gemini seems to return the string
@@ -135,7 +134,7 @@ class GlobalMessageHistory:
             Returns the AI-optimized history for the specified thread or newest thread.
         get_display_history(thread_id: Optional[ThreadID] = None) -> List[ChatCompletionMessageParam]:
             Returns the display-optimized history for the specified thread or newest thread.
-        append_message(ai_optimized_message: ChatCompletionMessageParam, display_message: ChatCompletionMessageParam, llm_provider: OpenAIProvider, thread_id: Optional[ThreadID] = None) -> None:
+        append_message(ai_optimized_message: ChatCompletionMessageParam, display_message: ChatCompletionMessageParam, llm_provider: ProviderManager, thread_id: Optional[ThreadID] = None) -> None:
             Appends messages to the specified thread (or newest thread) and generates a name if needed.
         truncate_histories(index: int, thread_id: Optional[ThreadID] = None) -> None:
             Truncates messages at the given index for the specified thread.
@@ -265,8 +264,7 @@ class GlobalMessageHistory:
         self, 
         ai_optimized_message: ChatCompletionMessageParam, 
         display_message: ChatCompletionMessageParam, 
-        model: str,
-        llm_provider: OpenAIProvider,
+        llm_provider: ProviderManager,
         thread_id: ThreadID
     ) -> None:
         """
@@ -305,7 +303,7 @@ class GlobalMessageHistory:
 
         # Outside the lock, await the name generation if needed
         if name_gen_input:
-            new_name = await generate_short_chat_name(str(name_gen_input[0]), str(name_gen_input[1]), model, llm_provider)
+            new_name = await generate_short_chat_name(str(name_gen_input[0]), str(name_gen_input[1]), llm_provider)
             with self._lock:
                 # Update the thread's name if still required
                 thread = self._chat_threads[thread_id]

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: